from torch import optim
from collections import deque
import json
import random
import numpy as np
import os

from parlai.core.agents import Agent
from parlai.core.build_data import modelzoo_path
from parlai.core.dict import DictionaryAgent
from parlai.core.utils import (
    AttrDict, argsort, padded_tensor, warn_once, round_sigfigs
)
from parlai.core.distributed_utils import is_primary_worker

try:
    import torch
except ImportError:
    raise ImportError('Need to install Pytorch: go to pytorch.org')

class Batch(AttrDict):
    """
    Batch is a namedtuple containing data being sent to an agent.

    This is the input type of the train_step and eval_step functions.
    Agents can override the batchify function to return an extended namedtuple
    with additional fields if they would like, though we recommend calling the
    parent function to set up these fields as a base.

    :param text_vec:
        bsz x seqlen tensor containing the parsed text data.

    :param text_lengths:
        list of length bsz containing the lengths of the text in same order as
        text_vec; necessary for pack_padded_sequence.

    :param label_vec:
        bsz x seqlen tensor containing the parsed label (one per batch row).

    :param label_lengths:
        list of length bsz containing the lengths of the labels in same order as
        label_vec.

    :param labels:
        list of length bsz containing the selected label for each batch row (some
        datasets have multiple labels per input example).

    :param valid_indices:
        list of length bsz containing the original indices of each example in the
        batch. we use these to map predictions back to their proper row, since e.g.
        we may sort examples by their length or some examples may be invalid.

    :param candidates:
        list of lists of text. outer list has size bsz, inner lists vary in size
        based on the number of candidates for each row in the batch.

    :param candidate_vecs:
        list of lists of tensors. outer list has size bsz, inner lists vary in size
        based on the number of candidates for each row in the batch.

    :param image:
        list of image features in the format specified by the --image-mode arg.

    :param observations:
        the original observations in the batched order
    """

    def __init__(self, text_vec=None, text_lengths=None,
                 label_vec=None, label_lengths=None, labels=None,
                 valid_indices=None,
                 candidates=None, candidate_vecs=None,
                 image=None, observations=None,
                 **kwargs):
        super().__init__(
            text_vec=text_vec, text_lengths=text_lengths,
            label_vec=label_vec, label_lengths=label_lengths, labels=labels,
            valid_indices=valid_indices,
            candidates=candidates, candidate_vecs=candidate_vecs,
            image=image, observations=observations,
            **kwargs)

class Output(AttrDict):
    """
    Output is an object containing agent predictions.

    This is the expected return type of the train_step and eval_step functions,
    though agents can choose to return None if they do not want to answer.

    :param List[str] text:
        list of strings of length bsz containing the predictions of the model

    :param List[List[str]] text_candidates:
        list of lists of length bsz containing ranked predictions of the model.
        each sub-list is an ordered ranking of strings, of variable length.
    """

    def __init__(self, text=None, text_candidates=None, **kwargs):
        super().__init__(text=text, text_candidates=text_candidates, **kwargs)

class History(object):
    """
    History handles tracking the dialogue history/state over the course of an
    episode.

    History may also be used to track the history of any field.

    :param field:
        field in the observation to track over the course of the episode
        (defaults to 'text')

    :param vec_type:
        specify a 'list' or 'deque' to save the history in this object

    :param maxlen:
        if `vec_type` is 'deque', this sets the maximum length of that object

    :param p1_token:
        token indicating 'person 1'; opt must have 'person_tokens' set to True
        for this to be added

    :param p1_token:
        token indicating 'person 2'; opt must have 'person_tokens' set to True
        for this to be added

    :param dict_agent:
        DictionaryAgent object for tokenizing the history
    """
    def __init__(self, opt, field='text', vec_type='deque', maxlen=None,
                 size=-1, p1_token='__p1__', p2_token='__p2__',
                 dict_agent=None):
        self.field = field
        self.dict = dict_agent
        self.delimiter = opt.get('delimiter', '\n')
        self.delimiter_tok = self.parse(self.delimiter)
        self.size = size
        self.split_on_newln = opt.get('split_lines', False)

        # set up history objects
        if vec_type != 'deque' and vec_type != 'list':
            raise RuntimeError(
                'Type {} is not supported for history'.format(vec_type)
            )
        self.vec_type = vec_type
        self.max_len = maxlen

        self.history_strings = []
        self.history_vecs = []

        # person token args
        self.add_person_tokens = opt.get('person_tokens', False)
        self.add_p1_after_newln = opt.get('add_p1_after_newln', False)
        self.p1_token = p1_token
        self.p2_token = p2_token

        # tracking when to clear history
        self.reset_on_next_update = False

    def parse(self, text):
        """Tokenize text with the given dictionary."""
        return self.dict.txt2vec(text)

    def reset(self):
        """Clear the history."""
        self.history_strings = []
        self.history_vecs = []

    def _update_strings(self, text):
        if self.size > 0:
            while len(self.history_strings) >= self.size:
                self.history_strings.pop(0)
        self.history_strings.append(text)

    def _update_vecs(self, text):
        if self.size > 0:
            while len(self.history_vecs) >= self.size:
                self.history_vecs.pop(0)
        self.history_vecs.append(self.parse(text))

    def update_history(self, obs, add_next=None):
        """
        Update the history with the given observation.

        :param add_next:
            string to append to history prior to updating it with the
            observation
        """
        if self.reset_on_next_update:
            # this is the first example in a new episode, clear the previous
            # history
            self.reset()
            self.reset_on_next_update = False

        if add_next is not None:
            if self.add_person_tokens:
                add_next = self._add_person_tokens(add_next, self.p2_token)
            # update history string
            self._update_strings(add_next)
            # update history vecs
            self._update_vecs(add_next)

        if self.field in obs and obs[self.field] is not None:
            if self.split_on_newln:
                next_texts = obs[self.field].split('\n')
            else:
                next_texts = [obs[self.field]]
            for text in next_texts:
                if self.add_person_tokens:
                    text = self._add_person_tokens(obs[self.field],
                                                   self.p1_token,
                                                   self.add_p1_after_newln)
                # update history string
                self._update_strings(text)
                # update history vecs
                self._update_vecs(text)

        if obs.get('episode_done', True):
            # end of this episode, clear the history when we see a new example
            self.reset_on_next_update = True

    def get_history_str(self):
        """Returns the string version of the history."""
        if len(self.history_strings) > 0:
            return self.delimiter.join(self.history_strings)
        return None

    def get_history_vec(self):
        """Returns a vectorized version of the history."""
        if len(self.history_vecs) == 0:
            return None

        if self.vec_type == 'deque':
            history = deque(maxlen=self.max_len)
            for vec in self.history_vecs[:-1]:
                history.extend(vec)
                history.extend(self.delimiter_tok)
            history.extend(self.history_vecs[-1])
        else:
            # vec type is a list
            history = []
            for vec in self.history_vecs[:-1]:
                history += vec
                history += self.delimiter_tok
            history += self.history_vecs[-1]

        return history

    def get_history_vec_list(self):
        """Returns a list of history vecs."""
        return self.history_vecs

    def _add_person_tokens(self, text, token, add_after_newln=False):
        if add_after_newln:
            split = text.split('\n')
            split[-1] = token + ' ' + split[-1]
            return '\n'.join(split)
        else:
            return token + ' ' + text

class TorchAgent(Agent):
    """
    A provided base agent for any model that wants to use Torch.

    Exists to make it easier to implement a new agent.
    Not necessary, but reduces duplicated code.

    Many methods are intended to be either used as is when the default is
    acceptable, or to be overriden and called with super(), with the extra
    functionality added to the initial result. See the method comment for
    recommended behavior.

    This agent serves as a common framework for all ParlAI models which want
    to use PyTorch.
    """

    P1_TOKEN = '__p1__'
    P2_TOKEN = '__p2__'

    @classmethod
    def optim_opts(self):
        """Fetch optimizer selection.

        By default, collects everything in torch.optim, as well as importing:
        - qhm / qhmadam if installed from github.com/facebookresearch/qhoptim

        Override this (and probably call super()) to add your own optimizers.
        """
        # first pull torch.optim in
        optims = {k.lower(): v for k, v in optim.__dict__.items()
                  if not k.startswith('__') and k[0].isupper()}

        try:
            # https://openreview.net/pdf?id=S1fUpoR5FQ
            from qhoptim.pyt import QHM, QHAdam
            optims['qhm'] = QHM
            optims['qhadam'] = QHAdam
        except ImportError:
            # no QHM installed
            pass

        return optims

    @staticmethod
    def dictionary_class():
        """
        Return the dictionary class that this agent expects to use.

        Can be overriden if a more complex dictionary is required.
        """
        return DictionaryAgent

    @classmethod
    def history_class(cls):
        """Return the history class that this agent expects to use.

        Can be overriden if a more complex history is required.
        """
        return History

    @classmethod
    def add_cmdline_args(cls, argparser):
        """Add the default commandline args we expect most agents to want."""
        agent = argparser.add_argument_group('TorchAgent Arguments')
        # pretrained embedding arguments
        agent.add_argument(
            '-emb', '--embedding-type', default='random',
            choices=['random', 'glove', 'glove-fixed', 'glove-twitter-fixed',
                     'fasttext', 'fasttext-fixed', 'fasttext_cc',
                     'fasttext_cc-fixed'],
            help='Choose between different strategies for initializing word '
                 'embeddings. Default is random, but can also preinitialize '
                 'from Glove or Fasttext. Preinitialized embeddings can also '
                 'be fixed so they are not updated during training.'
        )
        agent.add_argument(
            '-embp', '--embedding-projection', default='random',
            help='If pretrained embeddings have a different dimensionality '
                 'than your embedding size, strategy for projecting to the '
                 'correct size. If the dimensions are the same, this is '
                 'ignored unless you append "-force" to your choice.'
        )
        # optimizer arguments
        agent.add_argument(
            '-opt', '--optimizer', default='sgd', choices=cls.optim_opts(),
            help='Choose between pytorch optimizers. Any member of torch.optim'
                 ' should be valid.'
        )
        agent.add_argument(
            '-lr', '--learningrate', type=float, default=1,
            help='learning rate'
        )
        agent.add_argument(
            '-clip', '--gradient-clip', type=float, default=0.1,
            help='gradient clipping using l2 norm'
        )
        agent.add_argument(
            '-mom', '--momentum', default=0, type=float,
            help='if applicable, momentum value for optimizer.'
        )
        agent.add_argument(
            '--nesterov', default=True, type='bool',
            help='if applicable, whether to use nesterov momentum.'
        )
        agent.add_argument(
            '-nu', '--nus', default='0.7', type='floats',
            help='if applicable, nu value(s) for optimizer. can use a single '
                 'value like 0.7 or a comma-separated tuple like 0.7,1.0'
        )
        agent.add_argument(
            '-beta', '--betas', default='0.9,0.999', type='floats',
            help='if applicable, beta value(s) for optimizer. can use a single '
                 'value like 0.9 or a comma-separated tuple like 0.9,0.999'
        )
        # lr scheduler
        agent.add_argument(
            '--lr-scheduler', type=str, default='reduceonplateau',
            choices=['reduceonplateau', 'none', 'fixed', 'invsqrt'],
            help='Learning rate scheduler.'
        )
        agent.add_argument(
            '--lr-scheduler-patience', type=int, default=3,
            help='LR scheduler patience. In number of validation runs. If using '
                 'fixed scheduler, LR is decayed every <patience> validations.'
        )
        agent.add_argument(
            '--lr-scheduler-decay', type=float, default=0.5,
            help='Decay factor for LR scheduler, or how much LR is multiplied by '
                 'when it is lowered.'
        )
        agent.add_argument(
            '--warmup-updates', type=int, default=-1, hidden=True,
            help='Learning rate warmup period, in number of SGD updates. '
                 'Linearly scales up LR over period. Only enabled if > 0.'
        )
        agent.add_argument(
            '--warmup-rate', type=float, default=1e-4, hidden=True,
            help='Warmup learning rate *multiplier*. Initial LR is multiplied by '
                 'this value. Linearly adjusted up to 1.0 across --warmup-updates '
                 'steps.'
        )
        agent.add_argument(
            '--update-freq', type=int, default=-1, hidden=True,
            help='Accumulate gradients N times before performing an optimizer.step().'
        )
        # preprocessing arguments
        agent.add_argument(
            '-rc', '--rank-candidates', type='bool', default=False,
            help='Whether the model should parse candidates for ranking.'
        )
        agent.add_argument(
            '-tr', '--truncate', default=-1, type=int,
            help='Truncate input lengths to increase speed / use less memory.')
        agent.add_argument(
            '--text-truncate', type=int,
            help='Text input truncation length: if not specified, this will '
                 'default to `truncate`'
        )
        agent.add_argument(
            '--label-truncate', type=int,
            help='Label truncation length: if not specified, this will default '
                 'to `truncate`'
        )
        agent.add_argument(
            '-histsz', '--history-size', default=-1, type=int,
            help='Number of past dialog utterances to remember.'
        )
        agent.add_argument(
            '-pt', '--person-tokens', type='bool', default=False,
            help='add person tokens to history. adds __p1__ in front of input '
                 'text and __p2__ in front of past labels when available or '
                 'past utterances generated by the model. these are added to '
                 'the dictionary during initialization.'
        )
        agent.add_argument(
            '--split-lines', type='bool', default=False,
            help='split the dialogue history on newlines and save in separate '
                 'vectors'
        )
        agent.add_argument(
            '--use-reply', default='label', hidden=True,
            choices=['label', 'model', 'none'],
            help='Which previous replies to use as history. If label, use '
                 'gold dataset replies. If model, use model\'s own replies. '
                 'If none, do not track replies in history.'
        )
        agent.add_argument(
            '--add-p1-after-newln', type='bool', default=False, hidden=True,
            help='Add the other speaker token before the last newline in the '
                 'input instead of at the beginning of the input. this is '
                 'useful for tasks that include some kind of context before '
                 'the actual utterance (e.g. squad, babi, personachat).')
        agent.add_argument(
            '--delimiter', type=str, default='\n',
            help='Join history lines with this token, defaults to newline'
        )
        # GPU arguments
        # these gpu options are all mutually exclusive, and should error if the
        # user tries to present multiple of them
        gpugroup = agent.add_mutually_exclusive_group()
        gpugroup.add_argument(
            '-gpu', '--gpu', type=int, default=-1, help='which GPU to use'
        )
        gpugroup.add_argument(
            '--no-cuda', default=False, action='store_true', dest='no_cuda',
            help='disable GPUs even if available. otherwise, will use GPUs if '
                 'available on the device.'
        )

        cls.dictionary_class().add_cmdline_args(argparser)

    def __init__(self, opt, shared=None):
        """Initialize agent."""
        super().__init__(opt, shared)
        opt = self.opt
        if not shared:
            # intitialize any important structures from scratch
            self.replies = {}  # past replies
            self.dict = self.dictionary_class()(opt)
            if opt.get('person_tokens'):
                self.dict[self.P1_TOKEN] = 999999999
                self.dict[self.P2_TOKEN] = 999999998
        else:
            # copy initialized data from shared table
            self.opt = shared['opt']
            self.dict = shared['dict']
            if self.opt['batchsize'] == 1:
                # if we're not using batching (e.g. mturk), then replies really need
                # to stay separated
                self.replies = {}
            else:
                self.replies = shared['replies']

        if opt.get('numthreads', 1) > 1:
            torch.set_num_threads(1)

        # check for cuda
        self.use_cuda = not opt['no_cuda'] and torch.cuda.is_available()
        if self.use_cuda:
            if not shared:
                print('[ Using CUDA ]')
            if not shared and opt['gpu'] != -1:
                torch.cuda.set_device(opt['gpu'])

        # now set up any fields that all instances may need
        self.id = 'TorchAgent'  # child can override
        self.EMPTY = torch.LongTensor([])
        self.NULL_IDX = self.dict[self.dict.null_token]
        self.START_IDX = self.dict[self.dict.start_token]
        self.END_IDX = self.dict[self.dict.end_token]

        # for gradient acumulation
        self._number_grad_accum = 0
        # for the LR scheduler
        self._number_training_updates = 0
        # fixed random seed
        self.random = random.Random(42)
        # which row in the batch this instance is
        self.batch_idx = shared and shared.get('batchindex') or 0
        # can remember as few as zero utterances if desired
        self.histsz = opt['history_size']
        # truncate == 0 might give funny behavior
        self.truncate = opt['truncate'] if opt['truncate'] >= 0 else None
        text_truncate = opt.get('text_truncate') or opt['truncate']
        self.text_truncate = text_truncate if text_truncate >= 0 else None
        label_truncate = opt.get('label_truncate') or opt['truncate']
        self.label_truncate = label_truncate if label_truncate >= 0 else None
        # stores up to hist_utt past observations within current dialog
        self.history = self.history_class()(
            opt,
            maxlen=self.text_truncate,
            size=self.histsz,
            p1_token=self.P1_TOKEN,
            p2_token=self.P2_TOKEN,
            dict_agent=self.dict,
        )

        self.rank_candidates = opt['rank_candidates']
        self.add_person_tokens = opt.get('person_tokens', False)

    def _get_init_model(self, opt, shared):
        """Get model file to initialize with. If `init_model` exits, we will
        return the path to that file and maybe load dict file from that path.
        Otherwise, use `model_file.`

        :return:  path to load model from, whether we loaded from `init_model`
                  or not
        """
        init_model = None
        is_finetune = False
        if not shared:  # only do this on first setup
            # first check load path in case we need to override paths
            if opt.get('init_model') and os.path.isfile(opt['init_model']):
                # check first for 'init_model' for loading model from file
                init_model = opt['init_model']
                is_finetune = True
            if opt.get('model_file') and os.path.isfile(opt['model_file']):
                # next check for 'model_file', this would override init_model
                init_model = opt['model_file']
                is_finetune = False

            if init_model is not None:
                # if we are loading a model, should load its dict too
                if (os.path.isfile(init_model + '.dict') or
                        opt['dict_file'] is None):
                    opt['dict_file'] = init_model + '.dict'

        return init_model, is_finetune

    def init_optim(self, params, optim_states=None, saved_optim_type=None):
        """
        Initialize optimizer with model parameters.

        :param params:
            parameters from the model

        :param optim_states:
            optional argument providing states of optimizer to load

        :param saved_optim_type:
            type of optimizer being loaded, if changed will skip loading
            optimizer states
        """

        opt = self.opt

        # set up optimizer args
        lr = opt['learningrate']
        kwargs = {'lr': lr}
        if opt.get('momentum') > 0 and opt['optimizer'] in ['sgd', 'rmsprop', 'qhm']:
            # turn on momentum for optimizers that use it
            kwargs['momentum'] = opt['momentum']
            if opt['optimizer'] == 'sgd' and opt.get('nesterov', True):
                # for sgd, maybe nesterov
                kwargs['nesterov'] = opt.get('nesterov', True)
            elif opt['optimizer'] == 'qhm':
                # qhm needs a nu
                kwargs['nu'] = opt.get('nus', (0.7,))[0]
        elif opt['optimizer'] == 'adam':
            # turn on amsgrad for adam
            # amsgrad paper: https://openreview.net/forum?id=ryQu7f-RZ
            kwargs['amsgrad'] = True
        elif opt['optimizer'] == 'qhadam':
            # set nus for qhadam
            kwargs['nus'] = opt.get('nus', (0.7, 1.0))
        if opt['optimizer'] in ['adam', 'sparseadam', 'adamax', 'qhadam']:
            # set betas for optims that use it
            kwargs['betas'] = opt.get('betas', (0.9, 0.999))

        optim_class = self.optim_opts()[opt['optimizer']]
        self.optimizer = optim_class(params, **kwargs)
        if optim_states:
            if saved_optim_type != opt['optimizer']:
                print('WARNING: not loading optim state since optim class '
                      'changed.')
            else:
                try:
                    self.optimizer.load_state_dict(optim_states)
                except ValueError:
                    print('WARNING: not loading optim state since model '
                          'params changed.')
                if self.use_cuda:
                    for state in self.optimizer.state.values():
                        for k, v in state.items():
                            if isinstance(v, torch.Tensor):
                                state[k] = v.cuda()

    def build_lr_scheduler(self, states=None, hard_reset=False):
        """
        Create the learning rate scheduler, and assign it to self.scheduler.
        This scheduler will be updated upon a call to receive_metrics.

        May also create self.warmup_scheduler, if appropriate.

        :param state_dict states: Possible state_dict provided by model
            checkpoint, for restoring LR state

        :param bool hard_reset: If true, the LR scheduler should ignore the
            state dictionary.
        """
        if self.opt.get('warmup_updates', -1) > 0:
            def _warmup_lr(step):
                start = self.opt['warmup_rate']
                end = 1.0
                progress = min(1.0, step / self.opt['warmup_updates'])
                lr_mult = start + (end - start) * progress
                return lr_mult

            self.warmup_scheduler = optim.lr_scheduler.LambdaLR(
                self.optimizer,
                _warmup_lr,
            )
        else:
            self.warmup_scheduler = None

        patience = self.opt.get('lr_scheduler_patience', 3)
        decay = self.opt.get('lr_scheduler_decay', 0.5)

        if self.opt.get('lr_scheduler') == 'none':
            self.scheduler = None
        elif decay == 1.0:
            warn_once(
                "Your LR decay is set to 1.0. Assuming you meant you wanted "
                "to disable learning rate scheduling. Adjust --lr-scheduler-decay "
                "if this is not correct."
            )
            self.scheduler = None
        elif self.opt.get('lr_scheduler') == 'reduceonplateau':
            self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer,
                'min',
                factor=decay,
                patience=patience,
                verbose=True
            )
        elif self.opt.get('lr_scheduler') == 'fixed':
            self.scheduler = optim.lr_scheduler.StepLR(
                self.optimizer,
                patience,
                gamma=decay,
            )
        elif self.opt.get('lr_scheduler') == 'invsqrt':
            if self.opt.get('warmup_updates', -1) <= 0:
                raise ValueError(
                    '--lr-scheduler invsqrt requires setting --warmup-updates'
                )
            warmup_updates = self.opt['warmup_updates']
            decay_factor = np.sqrt(max(1, warmup_updates))

            def _invsqrt_lr(step):
                return decay_factor / np.sqrt(max(1, step))

            self.scheduler = optim.lr_scheduler.LambdaLR(
                self.optimizer,
                _invsqrt_lr,
            )
        else:
            raise ValueError(
                "Don't know what to do with lr_scheduler '{}'"
                .format(self.opt.get('lr_scheduler'))
            )

        # time to load LR state from the checkpoint, if possible.
        if states is None:
            # first make sure there are no null pointers
            states = {}

        if states and states.get('lr_scheduler_type') != self.opt['lr_scheduler']:
            # the LR scheduler changed, start things fresh
            warn_once("LR scheduler is different from saved. Starting fresh!")
            hard_reset = True

        if hard_reset:
            # We're not going to use the LR schedule, let's just exit
            return

        # do the actual loading (if possible)
        if 'number_training_updates' in states:
            self._number_training_updates = states['number_training_updates']
        if self.scheduler and 'lr_scheduler' in states:
            self.scheduler.load_state_dict(states['lr_scheduler'])
        if states.get('warmup_scheduler') and getattr(self, 'warmup_scheduler'):
            self.warmup_scheduler.load_state_dict(states['warmup_scheduler'])

    def report(self):
        metrics = {}
        # only report LR if we have a scheduler
        if hasattr(self, 'scheduler') and self.scheduler is not None:
            current_lr = round_sigfigs(self.optimizer.param_groups[0]['lr'], 4)
            metrics['lr'] = round_sigfigs(current_lr, 4)
        metrics['num_updates'] = self._number_training_updates
        return metrics

    def _is_lr_warming_up(self):
        """Checks if we're warming up the learning rate."""
        return (
            self.warmup_scheduler is not None and
            self._number_training_updates <= self.opt['warmup_updates']
        )

    def receive_metrics(self, metrics_dict):
        """
        Use the metrics to decide when to adjust LR schedule.

        This uses the loss as the validation metric if present, if not this
        function does nothing. Note that the model must be reporting loss for
        this to work.

        Override this to override the behavior.
        """
        if not hasattr(self, 'scheduler') or self.scheduler is None:
            return

        if self._is_lr_warming_up():
            # we're not done warming up, so don't start using validation
            # metrics to adjust schedule
            return

        if self.opt['lr_scheduler'] == 'none':
            # no scheduler, nothing to adjust here
            pass
        elif self.opt['lr_scheduler'] == 'reduceonplateau':
            if 'loss' not in metrics_dict:
                # nothing to step on, just skip
                warn_once("LR scheduler expected to see loss metric, but didn't.")
                return
            self.scheduler.step(metrics_dict['loss'])
        elif self.opt['lr_scheduler'] == 'fixed':
            self.scheduler.step()
        elif self.opt['lr_scheduler'] == 'invsqrt':
            # this is a training step lr scheduler, nothing to adjust in validation
            pass
        else:
            raise ValueError(
                "Don't know how to work with lr scheduler '{}'"
                .format(self.opt['lr_scheduler'])
            )

    def _get_embtype(self, emb_type):
        # set up preinitialized embeddings
        try:
            import torchtext.vocab as vocab
        except ImportError as ex:
            print('Please install torch text with `pip install torchtext`')
            raise ex
        pretrained_dim = 300
        if emb_type.startswith('glove'):
            if 'twitter' in emb_type:
                init = 'glove-twitter'
                name = 'twitter.27B'
                pretrained_dim = 200
            else:
                init = 'glove'
                name = '840B'
            embs = vocab.GloVe(
                name=name, dim=pretrained_dim,
                cache=modelzoo_path(self.opt.get('datapath'),
                                    'models:glove_vectors'))
        elif emb_type.startswith('fasttext_cc'):
            init = 'fasttext_cc'
            from parlai.zoo.fasttext_cc_vectors.build import url as fasttext_cc_url
            embs = vocab.Vectors(
                name='crawl-300d-2M.vec',
                url=fasttext_cc_url,
                cache=modelzoo_path(self.opt.get('datapath'),
                                    'models:fasttext_cc_vectors'))
        elif emb_type.startswith('fasttext'):
            init = 'fasttext'
            embs = vocab.FastText(
                language='en',
                cache=modelzoo_path(self.opt.get('datapath'),
                                    'models:fasttext_vectors'))
        else:
            raise RuntimeError('embedding type {} not implemented. check arg, '
                               'submit PR to this function, or override it.'
                               ''.format(emb_type))
        return embs, init

    def _project_vec(self, vec, target_dim, method='random'):
        """
        If needed, project vector to target dimensionality.

        Projection methods implemented are the following:

        random - random gaussian matrix multiplication of input vector

        :param vec:
            one-dimensional vector

        :param target_dim:
            dimension of returned vector

        :param method:
            projection method. will be used even if the dim is not changing if
            method ends in "-force".
        """
        pre_dim = vec.size(0)
        if pre_dim != target_dim or method.endswith('force'):
            if method.startswith('random'):
                # random projection
                if not hasattr(self, 'proj_rp'):
                    self.proj_rp = torch.Tensor(pre_dim, target_dim).normal_()
                    # rescale so we're not destroying norms too much
                    # http://scikit-learn.org/stable/modules/random_projection.html#gaussian-random-projection
                    self.proj_rp /= target_dim
                return torch.mm(vec.unsqueeze(0), self.proj_rp)
            else:
                # TODO: PCA
                # TODO: PCA + RP
                # TODO: copy
                raise RuntimeError('Projection method not implemented: {}'
                                   ''.format(method))
        else:
            return vec

    def _copy_embeddings(self, weight, emb_type, log=True):
        """
        Copy embeddings from the pretrained embeddings to the lookuptable.

        :param weight:
            weights of lookup table (nn.Embedding/nn.EmbeddingBag)

        :param emb_type:
            pretrained embedding type
        """
        if not is_primary_worker():
            # we're in distributed mode, copying embeddings in the workers
            # slows things down considerably
            return
        embs, name = self._get_embtype(emb_type)
        cnt = 0
        for w, i in self.dict.tok2ind.items():
            if w in embs.stoi:
                vec = self._project_vec(embs.vectors[embs.stoi[w]],
                                        weight.size(1))
                weight.data[i] = vec
                cnt += 1

        if log:
            print('Initialized embeddings for {} tokens ({}%) from {}.'
                  ''.format(cnt, round(cnt * 100 / len(self.dict), 1), name))

    def share(self):
        """
        Share fields from parent as well as useful objects in this class.

        Subclasses will likely want to share their model as well.
        """
        shared = super().share()
        shared['opt'] = self.opt
        shared['dict'] = self.dict
        shared['replies'] = self.replies
        return shared

    def _add_start_end_tokens(self, vec, add_start=False, add_end=False):
        if add_start:
            vec.insert(0, self.START_IDX)
        if add_end:
            vec.append(self.END_IDX)
        return vec

    def _v2t(self, vec):
        """Convert token indices to string of tokens."""
        new_vec = []
        if hasattr(vec, 'cpu'):
            vec = vec.cpu()
        for i in vec:
            if i == self.END_IDX:
                break
            new_vec.append(i)
        return self.dict.vec2txt(new_vec)

    def _vectorize_text(self, text, add_start=False, add_end=False,
                        truncate=None, truncate_left=True):
        """
        Return vector from text.

        :param text:
            String to vectorize.

        :param add_start:
            Add the start token to the front of the tensor.

        :param add_end:
            Add the end token to the end of the tensor.

        :param truncate:
            Truncate to this many tokens >= 0, or None.

        :param truncate_left:
            Truncate from the left side (keep the rightmost tokens). You
            probably want this True for inputs, False for targets.
        """
        vec = self.dict.txt2vec(text)
        vec = self._add_start_end_tokens(vec, add_start, add_end)
        vec = self._check_truncate(vec, truncate, truncate_left)
        tensor = torch.LongTensor(vec)
        return tensor

    def _check_truncate(self, vec, truncate, truncate_left=False):
        """Check that vector is truncated correctly."""
        if truncate is None:
            return vec
        if len(vec) <= truncate:
            return vec
        if truncate_left:
            return vec[-truncate:]
        else:
            return vec[:truncate]

    def _set_text_vec(self, obs, history, truncate):
        """
        Sets the 'text_vec' field in the observation.

        Useful to override to change vectorization behavior
        """
        if 'text' not in obs:
            return obs

        if 'text_vec' not in obs:
            # text vec is not precomputed, so we set it using the history
            obs['text'] = history.get_history_str()
            if obs['text'] is not None:
                obs['text_vec'] = history.get_history_vec()

        # check truncation
        if 'text_vec' in obs:
            obs['text_vec'] = torch.LongTensor(
                self._check_truncate(obs['text_vec'], truncate, True)
            )

        return obs

    def _set_label_vec(self, obs, add_start, add_end, truncate):
        """
        Sets the 'labels_vec' field in the observation.

        Useful to override to change vectorization behavior
        """
        # convert 'labels' or 'eval_labels' into vectors
        if 'labels' in obs:
            label_type = 'labels'
        elif 'eval_labels' in obs:
            label_type = 'eval_labels'
        else:
            label_type = None

        if label_type is None:
            return

        elif label_type + '_vec' in obs:
            # check truncation of pre-computed vector
            obs[label_type + '_vec'] = self._check_truncate(
                obs[label_type + '_vec'], truncate)
        else:
            # pick one label if there are multiple
            lbls = obs[label_type]
            label = lbls[0] if len(lbls) == 1 else self.random.choice(lbls)
            vec_label = self._vectorize_text(label, add_start, add_end,
                                             truncate, False)
            obs[label_type + '_vec'] = vec_label
            obs[label_type + '_choice'] = label

        return obs

    def _set_label_cands_vec(self, obs, add_start, add_end, truncate):
        """
        Sets the 'label_candidates_vec' field in the observation.

        Useful to override to change vectorization behavior
        """
        if 'label_candidates_vecs' in obs:
            if truncate is not None:
                # check truncation of pre-computed vectors
                vecs = obs['label_candidates_vecs']
                for i, c in enumerate(vecs):
                    vecs[i] = self._check_truncate(c, truncate)
        elif self.rank_candidates and obs.get('label_candidates'):
            obs['label_candidates'] = list(obs['label_candidates'])
            obs['label_candidates_vecs'] = [
                self._vectorize_text(c, add_start, add_end, truncate, False)
                for c in obs['label_candidates']]
        return obs

    def vectorize(self, obs, history, add_start=True, add_end=True,
                  text_truncate=None, label_truncate=None):
        """
        Make vectors out of observation fields and store in the observation.

        In particular, the 'text' and 'labels'/'eval_labels' fields are
        processed and a new field is added to the observation with the suffix
        '_vec'.

        If you want to use additional fields on your subclass, you can override
        this function, call super().vectorize(...) to process the text and
        labels, and then process the other fields in your subclass.

        :param obs:
            Single observation from observe function.

        :param add_start:
            default True, adds the start token to each label.

        :param add_end:
            default True, adds the end token to each label.

        :param text_truncate:
            default None, if set truncates text vectors to the specified
            length.

        :param label_truncate:
            default None, if set truncates label vectors to the specified
            length.

        :return:
            the input observation, with 'text_vec', 'label_vec', and
            'cands_vec' fields added.
        """
        self._set_text_vec(obs, history, text_truncate)
        self._set_label_vec(obs, add_start, add_end, label_truncate)
        self._set_label_cands_vec(obs, add_start, add_end, label_truncate)
        return obs

    def is_valid(self, obs):
        return 'text_vec' in obs or 'image' in obs

    def batchify(self, obs_batch, sort=False):
        """
        Create a batch of valid observations from an unchecked batch.

        A valid observation is one that passes the lambda provided to the
        function, which defaults to checking if the preprocessed 'text_vec'
        field is present which would have been set by this agent's 'vectorize'
        function.

        Returns a namedtuple Batch. See original definition above for in-depth
        explanation of each field.

        If you want to include additonal fields in the batch, you can subclass
        this function and return your own "Batch" namedtuple: copy the Batch
        namedtuple at the top of this class, and then add whatever additional
        fields that you want to be able to access. You can then call
        super().batchify(...) to set up the original fields and then set up the
        additional fields in your subclass and return that batch instead.

        :param obs_batch:
            List of vectorized observations

        :param sort:
            Default False, orders the observations by length of vectors. Set to
            true when using torch.nn.utils.rnn.pack_padded_sequence.  Uses the text
            vectors if available, otherwise uses the label vectors if available.
        """
        if len(obs_batch) == 0:
            return Batch()

        valid_obs = [(i, ex) for i, ex in enumerate(obs_batch) if
                     self.is_valid(ex)]

        if len(valid_obs) == 0:
            return Batch()

        valid_inds, exs = zip(*valid_obs)

        # TEXT
        xs, x_lens = None, None
        if any('text_vec' in ex for ex in exs):
            _xs = [ex.get('text_vec', self.EMPTY) for ex in exs]
            xs, x_lens = padded_tensor(_xs, self.NULL_IDX, self.use_cuda)
            if sort:
                sort = False  # now we won't sort on labels
                xs, x_lens, valid_inds, exs = argsort(
                    x_lens, xs, x_lens, valid_inds, exs, descending=True
                )

        # LABELS
        labels_avail = any('labels_vec' in ex for ex in exs)
        some_labels_avail = (labels_avail or
                             any('eval_labels_vec' in ex for ex in exs))

        ys, y_lens, labels = None, None, None
        if some_labels_avail:
            field = 'labels' if labels_avail else 'eval_labels'

            label_vecs = [ex.get(field + '_vec', self.EMPTY) for ex in exs]
            labels = [ex.get(field + '_choice') for ex in exs]
            y_lens = [y.shape[0] for y in label_vecs]

            ys, y_lens = padded_tensor(label_vecs, self.NULL_IDX, self.use_cuda)
            if sort and xs is None:
                ys, valid_inds, label_vecs, labels, y_lens = argsort(
                    y_lens, ys, valid_inds, label_vecs, labels, y_lens,
                    descending=True
                )

        # LABEL_CANDIDATES
        cands, cand_vecs = None, None
        if any('label_candidates_vecs' in ex for ex in exs):
            cands = [ex.get('label_candidates', None) for ex in exs]
            cand_vecs = [ex.get('label_candidates_vecs', None) for ex in exs]

        # IMAGE
        imgs = None
        if any('image' in ex for ex in exs):
            imgs = [ex.get('image', None) for ex in exs]

        return Batch(text_vec=xs, text_lengths=x_lens, label_vec=ys,
                     label_lengths=y_lens, labels=labels,
                     valid_indices=valid_inds, candidates=cands,
                     candidate_vecs=cand_vecs, image=imgs,
                     observations=exs)

    def match_batch(self, batch_reply, valid_inds, output=None):
        """
        Match sub-batch of predictions to the original batch indices.

        Batches may be only partially filled (i.e when completing the remainder
        at the end of the validation or test set), or we may want to sort by
        e.g the length of the input sequences if using pack_padded_sequence.

        This matches rows back with their original row in the batch for
        calculating metrics like accuracy.

        If output is None (model choosing not to provide any predictions), we
        will just return the batch of replies.

        Otherwise, output should be a parlai.core.torch_agent.Output object.
        This is a namedtuple, which can provide text predictions and/or
        text_candidates predictions. If you would like to map additional
        fields into the batch_reply, you can override this method as well as
        providing your own namedtuple with additional fields.

        :param batch_reply:
            Full-batchsize list of message dictionaries to put responses into.

        :param valid_inds:
            Original indices of the predictions.

        :param output:
            Output namedtuple which contains sub-batchsize list of text outputs
            from model. May be None (default) if model chooses not to answer.
            This method will check for ``text`` and ``text_candidates`` fields.
        """
        if output is None:
            return batch_reply
        if output.text is not None:
            for i, response in zip(valid_inds, output.text):
                batch_reply[i]['text'] = response
        if output.text_candidates is not None:
            for i, cands in zip(valid_inds, output.text_candidates):
                batch_reply[i]['text_candidates'] = cands
        return batch_reply

    def last_reply(self, use_reply='label'):
        """Retrieve the last reply from the model.

        If available, we use the true label instead of the model's prediction.

        By default, batch_act stores the batch of replies and this method
        will extract the reply of the current instance from the batch.

        :param use_label:
            default true, use the label when available instead of the model's
            generated response.
        """
        # if the last observation was the end of an episode,
        # then we shouldn't use it as history
        if (use_reply == 'none' or not self.observation or
                self.observation.get('episode_done', True)):
            return None

        if use_reply == 'label':
            # first look for the true label, if we aren't on a new episode
            label_key = ('labels' if 'labels' in self.observation else
                         'eval_labels' if 'eval_labels' in self.observation
                         else None)
            if label_key is not None:
                lbls = self.observation[label_key]
                last_reply = (lbls[0] if len(lbls) == 1
                              else self.random.choice(lbls))
                return last_reply

        # otherwise, we use the last reply the model generated
        batch_reply = self.replies.get('batch_reply')
        if batch_reply is not None:
            return batch_reply[self.batch_idx].get('text')
        return None

    def _save_history(self, observations, replies):
        """Save the model replies to the history."""
        # make sure data structure is set up
        if 'predictions' not in self.replies:
            self.replies['predictions'] = {}
        if 'episode_ends' not in self.replies:
            self.replies['episode_ends'] = {}
        # shorthand
        preds = self.replies['predictions']
        ends = self.replies['episode_ends']
        for i, obs in enumerate(observations):
            # iterate through batch, saving replies
            if i not in preds:
                preds[i] = []
            if ends.get(i):
                # check whether *last* example was the end of an episode
                preds[i].clear()
            ends[i] = obs.get('episode_done', True)
            preds[i].append(replies[i].get('text'))

    def reply_history(self):
        """
        Get the model's predicted reply history within this episode.

        :param batch:
            (default False) return the reply history for every row in the
            batch, otherwise will return just for this example.

        :return:
            list of lists of strings, each of the past model replies in in the
            current episode. will be None wherever model did not reply.
        """
        # make sure in batch order
        preds = sorted((b, p) for b, p in self.replies['predictions'].items())
        return [p for b, p in preds]

    def observe(self, observation):
        """Process incoming message in preparation for producing a response.

        This includes remembering the past history of the conversation.
        """
        reply = self.last_reply(use_reply=self.opt.get('use_reply', 'label'))
        # update the history using the observation
        self.history.update_history(observation, add_next=reply)
        self.observation = observation
        return self.vectorize(self.observation, self.history,
                              text_truncate=self.text_truncate,
                              label_truncate=self.label_truncate)

    def save(self, path=None):
        """
        Save model parameters to path (or default to model_file arg).

        Override this method for more specific saving.
        """
        path = self.opt.get('model_file', None) if path is None else path

        if path:
            states = {}
            if hasattr(self, 'model'):  # save model params
                if hasattr(self.model, 'module'):
                    # did we wrap in a DistributedDataParallel
                    states['model'] = self.model.module.state_dict()
                else:
                    states['model'] = self.model.state_dict()

            if hasattr(self, 'optimizer'):  # save optimizer params
                states['optimizer'] = self.optimizer.state_dict()
                states['optimizer_type'] = self.opt['optimizer']

            # lr scheduler
            if torch.__version__.startswith('0.'):
                warn_once(
                    "Must upgrade to Pytorch 1.0 to save the state of your "
                    "LR scheduler."
                )
            else:
                states['number_training_updates'] = self._number_training_updates
                if getattr(self, 'scheduler'):
                    states['lr_scheduler'] = self.scheduler.state_dict()
                    states['lr_scheduler_type'] = self.opt['lr_scheduler']
                if getattr(self, 'warmup_scheduler'):
                    states['warmup_scheduler'] = self.warmup_scheduler.state_dict()

            if states:  # anything found to save?
                with open(path, 'wb') as write:
                    torch.save(states, write)

                # save opt file
                with open(path + '.opt', 'w') as handle:
                    if hasattr(self, 'model_version'):
                        self.opt['model_version'] = self.model_version()
                    json.dump(self.opt, handle)
                    # for convenience of working with jq, make sure there's a newline
                    handle.write('\n')

    def load_state_dict(self, state_dict):
        """
        Load the state dict into model.

        This is easily overridable to facilitate transfer of state dicts.
        """
        self.model.load_state_dict(state_dict)

    def load(self, path):
        """
        Return opt and model states.

        Override this method for more specific loading.
        """
        states = torch.load(path, map_location=lambda cpu, _: cpu)
        if 'model' in states:
            self.load_state_dict(states['model'])
        if 'optimizer' in states and hasattr(self, 'optimizer'):
            self.optimizer.load_state_dict(states['optimizer'])
        return states

    def reset(self):
        """Clear internal states."""
        self.observation = {}
        self.history.reset()
        self.replies.clear()
        self.reset_metrics()

    def act(self):
        """Call batch_act with the singleton batch."""
        return self.batch_act([self.observation])[0]

    def batch_act(self, observations):
        """
        Process a batch of observations (batchsize list of message dicts).

        These observations have been preprocessed by the observe method.

        Subclasses can override this for special functionality, but if the
        default behaviors are fine then just override the ``train_step`` and
        ``eval_step`` methods instead. The former is called when labels are
        present in the observations batch; otherwise, the latter is called.
        """
        batch_size = len(observations)
        # initialize a list of replies with this agent's id
        batch_reply = [{'id': self.getID()} for _ in range(batch_size)]

        # check if there are any labels available, if so we will train on them
        is_training = any('labels' in obs for obs in observations)

        # create a batch from the vectors
        batch = self.batchify(observations)

        if is_training:
            output = self.train_step(batch)
        else:
            with torch.no_grad():
                # save memory and compute by disabling autograd.
                # use `with torch.enable_grad()` to gain back graidients.
                output = self.eval_step(batch)

        if output is None:
            self.replies['batch_reply'] = None
            return batch_reply

        self.match_batch(batch_reply, batch.valid_indices, output)
        self.replies['batch_reply'] = batch_reply
        self._save_history(observations, batch_reply)  # save model predictions

        return batch_reply

    def train_step(self, batch):
        """[Abstract] Process one batch with training labels."""
        raise NotImplementedError(
            'Abstract class: user must implement train_step'
        )

    def eval_step(self, batch):
        """[Abstract] Process one batch but do not train on it."""
        raise NotImplementedError(
            'Abstract class: user must implement eval_step'
        )

    def update_params(self):
        """
        Perform step of optimization, clipping gradients and adjusting LR
        schedule if needed. Gradient accumulation is also performed if agent
        is called with --update-freq.

        It is recommended (but not forced) that you call this in train_step.
        """
        update_freq = self.opt.get('update_freq', 1)
        if update_freq > 1:
            # we're doing gradient accumulation, so we don't only want to step
            # every N updates instead
            self._number_grad_accum = (self._number_grad_accum + 1) % update_freq
            if self._number_grad_accum != 0:
                return

        # keep track up number of steps, compute warmup factor
        self._number_training_updates += 1

        # compute warmup adjustment if needed
        if self.opt.get('warmup_updates', -1) > 0:
            if not hasattr(self, 'warmup_scheduler'):
                raise RuntimeError(
                    'Looks like you forgot to call build_lr_scheduler'
                )
            if self._is_lr_warming_up():
                self.warmup_scheduler.step(epoch=self._number_training_updates)

        if self.opt.get('lr_scheduler') == 'invsqrt' and not self._is_lr_warming_up():
            # training step scheduler
            self.scheduler.step(self._number_training_updates)

        if self.opt.get('gradient_clip', -1) > 0:
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), self.opt['gradient_clip']
            )

        self.optimizer.step()

    def zero_grad(self):
        """
        Zero out optimizer.

        It is recommended you call this in train_step. It automatically handles
        gradient accumulation if agent is called with --update-freq.
        """
        if self._number_grad_accum != 0:
            # if we're accumulating gradients, don't actually zero things out yet.
            return

        self.optimizer.zero_grad()


""" Converts DataTables JSON format to URCVT """

import json
import sys
from rcvformats.conversions.ut_without_transfers import UTWithoutTransfersConverter

class InvalidDataTableInput(Exception):
    """ Error messages caught here, that may not have been caught in JS """

# pylint: disable=too-few-public-methods
class ReadDataTableJSON():
    """ A class to read DataTable-serialized data """

    def __init__(self, jsonData):
        self.jsonData = jsonData
        self.numRounds = None

    def _parse_config(self):
        """ Parses the config from the JSON data passed from upload.html """
        config = {}
        config['contest'] = self.jsonData['configElectionTitle']
        config['date'] = self.jsonData['configElectionDate']
        config['threshold'] = self.jsonData['configThreshold']

        if config['contest'] == "":
            raise InvalidDataTableInput("Must give this election a title")

        if config['date'] == "":
            raise InvalidDataTableInput("Must give this election a date")

        if config['threshold'] == "":
            raise InvalidDataTableInput("Threshold cannot be empty")

        return config

    def _fill_rest_of_rounds_after_elected(
            self,
            results,
            candidateName,
            electedOnRound,
            finalVoteTotal):
        assert electedOnRound < self.numRounds

        # All future rounds keep the same vote total
        for futureRoundNum in range(electedOnRound + 1, self.numRounds):
            results[futureRoundNum]['tally'][candidateName] = finalVoteTotal

    def _handle_candidate(self, results, candidateData, candidateName):
        """
        Fills in the data for one candidate across all rounds
        Returns true if was just elected
        """
        # Iterate over each row for this candidate
        wasElectedLastRound = False
        for roundNum, roundData in enumerate(candidateData):
            # Gather data from table
            numVotes = roundData['# Votes']
            status = roundData['Status']

            try:
                numVotes = float(numVotes)
            except (TypeError, ValueError) as exc:
                raise InvalidDataTableInput(f"On Round {roundNum+1}, \"{candidateName}\" " +
                                            "has an invalid number of votes") from exc

            if numVotes < 0:
                raise InvalidDataTableInput("All vote counts must be positive")

            # Already elected last round - fill out remaining rounds, then break
            if wasElectedLastRound:
                if numVotes is None or numVotes == 0:
                    # No surplus transfer: use the last-round's value
                    finalVoteTotal = results[roundNum - 1]['tally'][candidateName]
                else:
                    # Has surplus transfer
                    finalVoteTotal = float(numVotes)

                self._fill_rest_of_rounds_after_elected(
                    results, candidateName, roundNum - 1, finalVoteTotal)
                break

            # The normal case: just append the tally
            results[roundNum]['tally'][candidateName] = numVotes

            # TallyResults only exist for eliminated and elected candidates
            if status == 'Eliminated':
                results[roundNum]['tallyResults'].append({'eliminated': candidateName})
                break
            if status == 'Elected':
                results[roundNum]['tallyResults'].append({'elected': candidateName})
                wasElectedLastRound = True
            else:
                assert status == 'Active'

    def _parse_results(self):
        """ Parses the results from the dataEntry serialization """
        if 'dataEntry' not in self.jsonData:
            raise InvalidDataTableInput("Data Tables isn't even shown. This should never happen.")
        dataTableData = json.loads(self.jsonData['dataEntry'])

        self.numRounds = len(dataTableData['data'][0])
        results = [{"round": i + 1, "tally": {}, "tallyResults": []} for i in range(self.numRounds)]

        names = dataTableData['rowNames']
        if len(set(names)) != len(names):
            raise InvalidDataTableInput("All candidate names must be unique")
        if any(n == "" or n is None for n in names):
            raise InvalidDataTableInput("All candidates must have names")

        # Iterate over the data, which goes candidate-by-candidate (whereas URCVT data
        # goes round-by-round), so we will transpose it
        for candidateNum, candidateData in enumerate(dataTableData['data']):
            candidateName = names[candidateNum]
            self._handle_candidate(results, candidateData, candidateName)

        return results

    def convert_to_urcvt(self):
        """
        Worker function: converts all data passed from upload.html, including the config
        (from RCVis' fields) and the results (from the datatables' serialization)
        """
        config = self._parse_config()
        results = self._parse_results()
        withoutTransfers = {'config': config, 'results': results}

        converter = UTWithoutTransfersConverter()
        try:
            withTransfers = converter.convert_to_ut(withoutTransfers)
        except Exception as exc:
            # Try to find a reason
            totalVotesLastRound = sys.maxsize
            for roundNum, result in enumerate(results):
                totalVotesThisRound = sum(t for t in result['tally'].values())
                if totalVotesThisRound > totalVotesLastRound:
                    raise InvalidDataTableInput(
                        f"Round {roundNum+1} has {totalVotesThisRound} votes, which is more " +
                        f"than the previous round, which only has {totalVotesLastRound}. " +
                        "This should not be possible.") from exc
                totalVotesLastRound = totalVotesThisRound

            # No reason found
            raise InvalidDataTableInput("Could not add transfers") from exc

        return withTransfers


import gym
import time
env = gym.make('MountainCarContinuous-v0')
    #       Description:
    #     The agent (a car) is started at the bottom of a valley. For any given
    #     state the agent may choose to accelerate to the left, right or cease
    #     any acceleration.
for i_episode in range(5):
    cumreward = []

    observation = env.reset()
    print("The initial observation is {}".format(observation))
    # Starting State:
    #      The position of the car is assigned a uniform random value in
    #      [-0.6 , -0.4].
    #      The starting velocity of the car is always assigned to 0.

    for t in range(100):
        env.render(mode = 'rgb_array')
        action = env.action_space.sample()
        observation, reward, done, info = env.step(action)
        time.sleep(0.01)
        print("Observation {} ".format(observation))
        # Observation:
        #     Type: Box(2)
        #     Num    Observation               Min            Max
        #     0      Car Position              -1.2           0.6
        #     1      Car Velocity              -0.07          0.07

        print("Action {} ".format(action))
        # Actions:
        #     Type: Box(1)
        #     Num    Action                    Min            Max
        #     0      the power coef            -1.0           1.0
        #     Note: actual driving force is calculated by multiplying the power coef by power (0.0015)

        print("Reward {} ".format(reward))
        # Reward:
        #      Reward of 100 is awarded if the agent reached the flag (position = 0.45) on top of the mountain.
        #      Reward is decrease based on amount of energy consumed each step.
        cumreward.append(reward)

        if done:
            print("Cumulatitive reward {} ".format(sum(cumreward)))
            print("Episode finished after {} timesteps".format(t+1))
            break
        # Episode Termination:
        #      The car position is more than 0.45
        #      Episode length is greater than 200

print(reward,type(reward))
print(action,type(action))
print(observation,type(observation))

env.close()


import torch

from IBA.pytorch import IBA, tensor_to_np_img
from IBA.utils import plot_saliency_map

from reader.loader import get_test_loader, get_train_loader
from nets import model
import configs
import matplotlib.pyplot as plt
import os
from os.path import join, exists
import argparse
import cv2
import numpy as np
import argparse

def show_cam_on_image(img, mask, path="./e/cam.jpg"):
    #print(path)
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img)
    cam = cam / np.max(cam)
    cv2.imwrite(path, np.uint8(255 * cam))

def get_layer(net, arch, layer):
    if arch.startswith('dense'):
        if layer == 'denseblock1':
            feat_layer = net.model.features[4]
        elif layer == 'denseblock2':
            feat_layer = net.model.features[6]
        elif layer == 'denseblock3':
            feat_layer = net.model.features[8]
        elif layer == 'denseblock4':
            feat_layer = net.model.features[10]
        else:
            raise Exception('Unknown layer name {} for arch {}'.format(layer, arch))
    elif arch.startswith('nest') or arch.startswith('resnet'):
        if layer == 'layer1':
            feat_layer = net.model.layer1
        elif layer == 'layer2':
            feat_layer = net.model.layer2
        elif layer == 'layer3':
            feat_layer = net.model.layer3
        elif layer == 'layer4':
            feat_layer = net.model.layer4
        else:
            raise Exception('Unknown layer name {} for arch {}'.format(layer, arch))
    else:
        raise Exception('Unknown arch name')

    return feat_layer

def run(opt, layer, weight_postfix, n_samples, vis_out_dir, orig_out_dir):
    # Data Loader
    eval_loader = get_test_loader(**{**opt['dataset'], 'batch_size':1})
    train_loader = get_train_loader(**opt['dataset'])

    # Network configuration
    netconf = {k:v for k, v in opt['netA'].items() if k in ['dev', 'model']}
    if 'istrained' in netconf['model']: netconf['model']['istrained'] = False
    netA = model.build(name='netA', **netconf, resume=join(opt['modeldir'], 'chk_{}.pth.tar'.format(weight_postfix))).eval()
    print(netA)

    # Feature Layer Selection
    feat_layer = get_layer(netA, opt['netA']['model']['arch'], layer)
    iba = IBA(feat_layer)

    # IBA stats calculation
    iba.reset_estimate()
    iba.estimate(netA.model, train_loader, device=torch.device(netA.dev), n_samples=n_samples, progbar=True)

    for i, batch in enumerate(eval_loader):
        x, label = batch[0].to(netA.dev), batch[1].to(netA.dev)

        model_loss_closure = lambda x: -torch.log_softmax(netA(x), dim=1)[:, label].mean()
        heatmap = iba.analyze(x, model_loss_closure)

        # Visualization
        x = x.cpu()[0,:,:,:].permute(1,2,0).numpy()
        x = eval_loader.dataset.norm.scale_to_01(x)
        show_cam_on_image(x, heatmap, join(vis_out_dir, str(i)+ '_label_'+ str(label.item()) +'.png'))
        if orig_out_dir is not None:
            cv2.imwrite(join(orig_out_dir, str(i)+ '_label_'+ str(label.item()) +'.png'), np.uint8(255 * x))

if __name__ == "__main__":
    # Argumnent parsing
    parser = argparse.ArgumentParser()
    parser.add_argument('-e', type=str, required=True, help='Experiment name')
    parser.add_argument('-m', type=str, default='desktop', help='Machine name')
    parser.add_argument('-n', type=int, default=1000, help='# of samples for stat calculation')
    parser.add_argument('-w', type=str, default='best_v_accuracy', help='Postfix of the weight file')
    parser.add_argument('--layer', type=str, required=True, help='Feature layer name')
    parser.add_argument('--org', action='store_true', help='Save input images')
    args = parser.parse_args()

    # Configuration
    opt = configs.parse(args.m, args.e)

    # Create output directories
    vis_out_dir = join(opt['outdir'], 'vis', '[{}]_{}_iba'.format(args.w, args.layer))
    orig_out_dir = join(opt['outdir'], 'vis', 'org')
    if not exists(vis_out_dir): os.makedirs(vis_out_dir)
    if not exists(orig_out_dir) and args.org: os.makedirs(orig_out_dir)
    if not args.org: orig_out_dir = None
    run(opt, layer=args.layer, weight_postfix=args.w, n_samples=args.n, vis_out_dir=vis_out_dir, orig_out_dir=orig_out_dir)


from yahoo_finance import Share
import csv
import sqlite3
import os.path

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
db_path = os.path.join(BASE_DIR, "fortune500.db")

def get_stock_symbols(file_name):
    symbols = []
    with open(file_name, 'r') as file:
        file_reader = csv.reader(file, delimiter=',', quotechar='|')
        for row in file_reader:
            if len(row[0]) < 6:
                symbols.append((row[0], row[1]))
    return symbols

def get_stock_info(stock_symbol):
    _stock_info = dict()
    stock = Share(stock_symbol)
    _stock_info['price'] = stock.get_price()
    _stock_info['market_cap'] = stock.get_market_cap()
    if _stock_info['market_cap'] is not None:
        _stock_info['market_cap'] = _stock_info['market_cap'].replace("B", "")
    _stock_info['price_earnings'] = stock.get_price_earnings_ratio()
    return _stock_info

def get_stock_historical(stock_symbol, date_one, date_two):
    stock_info = dict()
    stock = Share(stock_symbol)
    stock_hist = stock.get_historical(date_one, date_two)
    return stock_hist

def get_all_stocks():
    symbols = get_stock_symbols('Fortune500Companies.csv')
    stock_info = dict()
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    for symbol, name in symbols:
        stock_info[symbol] = get_stock_info(symbol)
        stock_info[symbol]['name'] = name
        cursor.execute("""INSERT OR REPLACE INTO basic VALUES(?, ?, ?, ?, ?)""", (symbol, stock_info[symbol]['name'], stock_info[symbol]["price"], stock_info[symbol]["market_cap"], stock_info[symbol]["price_earnings"]))
    conn.commit()
    conn.close()

def save_stocks_to_file(open_file_name, save_file_name, start_date, end_date):
    symbols = get_stock_symbols(open_file_name)
    stock_info = dict()
    with open(save_file_name, 'w') as save_file:
        writer = csv.writer(save_file, delimiter=',', quotechar='|')
        writer.writerow(["date", "name", "symbol", "high", "low", "close", "adj close", "volume", "high market cap",
                         "low market cap", "close market cap", "adj close market cap"])
        for symbol, name in symbols:
            try:
                stock_info = get_stock_historical(symbol, start_date, end_date)
                for stock in stock_info:
                    writer.writerow([stock['Date'],
                                     name,
                                     stock['Symbol'],
                                     stock['High'],
                                     stock['Low'],
                                     stock['Close'],
                                     stock['Adj_Close'],
                                     stock['Volume'],
                                     float(stock['High']) * float(stock['Volume']),
                                     float(stock['Low']) * float(stock['Volume']),
                                     float(stock['Close']) * float(stock['Volume']),
                                     float(stock['Adj_Close']) * float(stock['Volume'])])
            except Exception:
                print("There was a problem but we ignored it")

# get_all_stocks()
save_stocks_to_file("Fortune500Companies.csv", "Stocks.csv", '2008-01-01', '2017-3-15')


# -*- coding: utf-8 -*-

"""
[CN]

该 Script Filter 的功能是根据 key 随机生成一个 1 ~ 1000 之间的 value. 这个 value 的值
将会被缓存 5 秒. 5 秒内查询同一个 key 的结果将会是一样的. 该例子用来展示如何使用 time to live
缓存.

在 Alfred Workflow 的 Canvas 界面中 Script Filter 的设置如下:

- Keyword: afwf-example-memorize-cache, Argument Required
- Language: /bin/bash
- Script: python main.py 'memorize_cache {query}'
"""

import random
import afwf
import attr

from ..cache import cache

@attr.define
class Handler(afwf.Handler):
    @cache.memoize(tag="memorize_cache", expire=5)
    def main(self, key: str) -> afwf.ScriptFilter:
        sf = afwf.ScriptFilter()
        value = random.randint(1, 1000)
        item = afwf.Item(
            title=f"value is {value}",
        )
        sf.items.append(item)
        return sf

    def parse_query(self, query: str):
        return dict(
            key=query,
        )

handler = Handler(id="memorize_cache")


#-*- coding:utf-8 -*-

import webdispatch

from webob import dec

from .controller import main, chat

def hello(environ, start_response):

    start_response('200 ok', [])

    return ['hello, world']

def make_application(skype, options):

    dispatcher = webdispatch.URLDispatcher()

    dispatcher.add_url('main', '/', dec.wsgify(main.Main(skype)))
    dispatcher.add_url('chatlist', '/chats/', dec.wsgify(chat.ChatList(skype)))
    dispatcher.add_url('chat', '/chats/*', dec.wsgify(chat.Chat(skype)))

    return dispatcher


from subprocess import run
import wave
from pydub import AudioSegment
import tensorflow as tf
import downloader
tf.config.set_visible_devices([], 'GPU')
from final_model import final_pred
from ffmpy import FFmpeg
import os

def trim(dir):
    wav_file = wave.open(dir, 'r')
    sample_rate = wav_file.getframerate()
    num_channels = wav_file.getnchannels()
    audio_data = AudioSegment.from_wav(dir)
    trimmed_data = audio_data[0:180000]
    trimmed_data.export(dir)
    wav_file.close()

def get_data(name):
    Path = downloader.get_path()
    print(Path)
    input_file = Path
    output_file = "# save file location\{}.wav".format(name)
    print(output_file)
    run(["ffmpeg", "-i", input_file, "-vn", "-acodec", "pcm_s16le", "-ar", "44100", "-ac", "2", output_file])
    print("----------")
    trim(output_file)

    os.remove(Path)
    return output_file

def main(title,artist):
    try:
        request = title + " " + artist
        downloader.main(request)
        path = get_data(title)
        print(path)
        return final_pred(path)
    except:
        print('Cannot Access', title)

pred = main("fix you","coldplay")
print(pred)


import argparse
import ast

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('rulesfile')
    args = parser.parse_args()

    ops_table = {}
    with open(args.rulesfile) as f:
        for line in f.readlines():
            if 'applied' in line or '--' in line:
                continue
            ops = ast.literal_eval(line.strip())
            for op in ops:
                if op in ops_table:
                    ops_table[op] += 1
                else:
                    ops_table[op] = 1

    # Write out the table
    with open(args.rulesfile + '.table', 'w') as f:
        pairs = []
        for key in ops_table.keys():
            pairs.append((ops_table[key], key))
        spairs = []
        for (n, v) in sorted(pairs):
            spairs.append(str(n) + ", " + str(v))
        f.write('\n'.join(spairs))


"""peaksters URL Configuration

The `urlpatterns` list routes URLs to views. For more information please see:
    https://docs.djangoproject.com/en/3.2/topics/http/urls/
Examples:
Function views
    1. Add an import:  from my_app import views
    2. Add a URL to urlpatterns:  path('', views.home, name='home')
Class-based views
    1. Add an import:  from other_app.views import Home
    2. Add a URL to urlpatterns:  path('', Home.as_view(), name='home')
Including another URLconf
    1. Import the include() function: from django.urls import include, path
    2. Add a URL to urlpatterns:  path('blog/', include('blog.urls'))
"""
from django.contrib import admin
from django.urls import path
from django.conf.urls import url
from books import views as books_view
from teams import views as teams_view

urlpatterns = [
    path('admin/', admin.site.urls),
    url(r'^home/$', books_view.home),
    url('previous-matches/$', books_view.previous_matches),
    url('detail/(?P<slug>[\w-]+)/$', books_view.details, name="detail"),
    url('detail/previous/(?P<slug>[\w-]+)/$', books_view.details_previous, name="detail_previous"),
    url('login/$',books_view.login_view),
    url('congratulations/$', books_view.congrats),
    url('profile/$', books_view.profile),
    url('logout/$', books_view.logout_view),
    url('rewards/handout/$' , books_view.entry_rewards),
    url('rewards/handout/(?P<bet_id>\w{0,50})/$', books_view.entry_detail, name="entry_detail"),

    # Teams Views
   # url('tournaments/$', teams_view.tournament),
   # url('tournaments/(?P<index>\w{0,50})/$', teams_view.tournament_detail, name="tournament_detail"),
    url('teams/$', teams_view.team),
    url('teams/(?P<index1>\w{0,50})/$', teams_view.team_detail, name="team_detail"),
    #url('blinders/codm/pointstable/$', teams_view.points_table , name="points_table_codm"),
]


import librosa
import numpy
import pyaudio
import time

# Define global variables.
CHANNELS = 1
RATE = 44100
FRAMES_PER_BUFFER = 1000
N_FFT = 4096
SCREEN_WIDTH = 178
ENERGY_THRESHOLD = 0.4

#frequency range of log-spectrogram.
F_LO = librosa.note_to_hz('C2')
F_HI = librosa.note_to_hz('C9')
M = librosa.filters.mel(RATE, N_FFT, SCREEN_WIDTH, fmin=F_LO, fmax=F_HI)

p = pyaudio.PyAudio()

def generate_string_from_audio(audio_data):
    # Compute real FFT.
    x_fft = numpy.fft.rfft(audio_data, n=N_FFT)

    # Compute mel spectrum.
    melspectrum = M.dot(abs(x_fft))

    # Initialize output characters to display.
    char_list = [' ']*SCREEN_WIDTH

    for i in range(SCREEN_WIDTH):

        if melspectrum[i] > ENERGY_THRESHOLD:
            char_list[i] = '*'

        elif i % 30 == 29:
            char_list[i] = '|'

    return ''.join(char_list)

def callback(in_data, frame_count, time_info, status):
    audio_data = numpy.fromstring(in_data, dtype=numpy.float32)
    print( generate_string_from_audio(audio_data) )
    return (in_data, pyaudio.paContinue)

stream = p.open(format=pyaudio.paFloat32,
                channels=CHANNELS,
                rate=RATE,
                input=True,   # Do record input.
                output=False, # Do not play back output.
                frames_per_buffer=FRAMES_PER_BUFFER,
                stream_callback=callback)

stream.start_stream()

while stream.is_active():
    time.sleep(0.100)

stream.stop_stream()
stream.close()

p.terminate()


# -*- coding: utf-8 -*-
"""
Created on Wed Mar  9 11:50:48 2022

@author: iant

CHECK LOIC VS GODDARD SPECTRAL CALIBRATION AND FIND WHICH ONE MATCHES DATA THE BEST
"""
import re
import numpy as np
import matplotlib.pyplot as plt
import h5py
import os
import pandas as pd

from matplotlib.backends.backend_pdf import PdfPages

# from scipy.optimize import curve_fit
# from scipy.optimize import OptimizeWarning

from tools.spectra.molecular_spectrum_so import get_molecular_hr
from tools.spectra.baseline_als import baseline_als
# from tools.spectra.fit_gaussian_absorption import fit_gaussian_absorption

from tools.plotting.colours import get_colours

from tools.general.get_minima_maxima import get_local_minima
from tools.general.get_nearest_index import get_nearest_index

from tools.file.hdf5_functions import make_filelist
from tools.file.paths import paths

from tools.file.write_log import write_log

order_dict = {
    119:{"molecule":"CO2", "smin":1.0e-30, "n_lines":20},
    121:{"molecule":"CO2", "smin":1.0e-30, "n_lines":20},
    134:{"molecule":"H2O", "smin":1.0e-28, "n_lines":9},
    136:{"molecule":"H2O", "smin":1.0e-28, "n_lines":7},
    149:{"molecule":"CO2", "smin":1.0e-28, "n_lines":17},
    165:{"molecule":"CO2", "smin":1.0e-24, "n_lines":20},
    168:{"molecule":"H2O", "smin":1.0e-33, "n_lines":5},
    169:{"molecule":"H2O", "smin":1.0e-33, "n_lines":6},
    186:{"molecule":"CO", "smin":1.0e-33, "n_lines":8},
    189:{"molecule":"CO", "smin":1.0e-33, "n_lines":9},
    190:{"molecule":"CO", "smin":1.0e-33, "n_lines":10},
    191:{"molecule":"CO", "smin":1.0e-33, "n_lines":10},

}

chosen_bin = 3
chosen_y_mean = 0.5

temperature_grid_spacing = 0.1

PLOT_SPECTRA = False

def gd21_waven(order, inter_temp, px_in):
    """spectral calibration Goddard Oct 21"""
    cfpixel = [1.75128E-08, 5.55953E-04, 2.24734E+01]             # Blaze free-spectral-range (FSR) [cm-1 from pixel]
    ncoeff  = [-2.44383699e-07, -2.30708836e-05, -1.90001923e-04] # Relative frequency shift coefficients [shift/frequency from Celsius]

    xdat  = np.polyval(cfpixel, px_in) * order
    xdat += xdat  *np.polyval(ncoeff, inter_temp)
    return xdat

def lt21_waven(order, inter_temp, px_in):
    """spectral calibration Loic Feb 22"""

    cfpixel = [3.32e-8, 5.480e-4, 22.4701]

    p0_nu = -0.8276 * inter_temp #px/°C * T(interpolated)
    px_temp = px_in + p0_nu
    xdat  = np.polyval(cfpixel, px_temp) * order

    return xdat

# px = np.arange(n_px)
# plt.figure()
# for order in range(119, 195, 10):
#     for inter_temp in range(-10, 10, 2):
#         nu_p_gd = gd21_waven(order, inter_temp, px)
#         nu_p_lt = lt21_waven(order, inter_temp, px)

#         delta = nu_p_gd - nu_p_lt
#         plt.plot(delta, label="%i %0.1fC" %(order, inter_temp))
#         plt.legend()
# stop()

# def fit_gaussian_absorption(x_in, y_in, error=False, hr_num=500):
#     """fit inverted gaussian to absorption band.
#     Normalise continuum to 1 first"""

#     def func(x, a, b, c, d):
#         return 1.0 - a * np.exp(-((x - b)/c)**2.0) + d

#     x_mean = np.mean(x_in)
#     x_centred = x_in - x_mean
#     try:
#         popt, pcov = curve_fit(func, x_centred, y_in, p0=[0.1, 0.02, 0.25, 0.0])
#     except OptimizeWarning :
#         if error:
#             return [], [], [], []
#         else:
#             return [], [], []
#     except RuntimeError:
#         if error:
#             return [], [], [], []
#         else:
#             return [], [], []

#     if error:
#         y_fit = func(x_centred, *popt)
#         chi_squared = np.sum(((y_in - y_fit) / y_fit)**2) #divide by yfit to normalise large and small absorption bands

#     x_hr = np.linspace(x_in[0], x_in[-1], num=hr_num)
#     y_hr = func(x_hr - x_mean, *popt)

#     min_index = (np.abs(y_hr - np.min(y_hr))).argmin()

#     x_min_position = x_hr[min_index]

#     if error:
#         return x_hr, y_hr, x_min_position, chi_squared
#     else:
#         return x_hr, y_hr, x_min_position

def fit_polynomial(x_in, y_in, deg=2, error=False, hr_num=100):

    x_mean = np.mean(x_in)
    x_centred = x_in - x_mean
    polyfit = np.polyfit(x_centred, y_in, deg)

    y_fit = np.polyval(polyfit, x_centred)

    if error:
        chi_squared = np.sum(((y_in - y_fit) / y_fit)**2) #divide by yfit to normalise large and small absorption bands

    x_hr = np.linspace(x_in[0], x_in[-1], num=hr_num)
    y_hr = np.polyval(polyfit, x_hr - x_mean)

    min_index = (np.abs(y_hr - np.min(y_hr))).argmin()

    x_min_position = x_hr[min_index]

    if error:
        return x_hr, y_hr, x_min_position, chi_squared
    else:
        return x_hr, y_hr, x_min_position

from matplotlib.colors import LinearSegmentedColormap

redgreyblue = LinearSegmentedColormap.from_list('redgreyblue', (
    (0.000, (0.000, 0.325, 0.984)),
    (0.480, (0.620, 0.627, 0.863)),
    (0.500, (0.596, 0.588, 0.588)),
    (0.520, (0.863, 0.620, 0.624)),
    (1.000, (1.000, 0.000, 0.047))))

def make_spectral_comparison_log(order_dict, chosen_bin, chosen_y_mean):

    write_log("Filename\tOrder\tTemperature\tAbsorption nu\tAbsorption pixel\tGD nu\tLT nu\tGD delta nu\tLT delta nu\tFitted depth\tFitted max")

    for order, order_info in order_dict.items():

        molecule = order_info["molecule"]
        smin = order_info["smin"]
        n_lines = order_info["n_lines"]

        n_px = 320.0 #consider whole detector
        file_level = "hdf5_level_1p0a"

        # regex = re.compile("20(18|21)04.._......_.*_SO_A_[IE]_%i" %order)
        # regex = re.compile("20(18|19|20|21)04.._......_.*_SO_A_[IE]_%i" %order)
        # regex = re.compile("20(18|19|20|21)0[4567].._......_.*_SO_A_[IE]_%i" %order)
        regex = re.compile("20(18|19|20|21|22)...._......_.*_SO_A_[IE]_%i" %order)

        hdf5_files, hdf5_filenames, _ = make_filelist(regex, file_level, open_files=False, silent=True)

        colours = get_colours(len(hdf5_filenames))

        if PLOT_SPECTRA:
            fig, (ax1,ax2) = plt.subplots(nrows=2, ncols=1, figsize=(18,10), sharex=True, constrained_layout=True)
            fig.suptitle("Order %i" %order)

        for file_index, (hdf5_file_path, hdf5_filename) in enumerate(zip(hdf5_files, hdf5_filenames)):

            hdf5_file = h5py.File(hdf5_file_path, "r")

            colour = colours[file_index]

            y_all = hdf5_file["Science/Y"][:, :]
            y_mean = np.mean(y_all, axis=1)

            if "InterpolatedTemperature" not in hdf5_file["Channel"].keys():
                print("%s: Error temperatures not in file" %hdf5_filename)
                continue

            temperatures = hdf5_file["Channel/InterpolatedTemperature"][...]

            bins = hdf5_file["Science/Bins"][:, 0]
            unique_bins = sorted(list(set(bins)))

            bin_indices = np.where(bins == unique_bins[chosen_bin])[0]

            y_bin = y_all[bin_indices, :]
            y_mean_bin = y_mean[bin_indices]

            index = get_nearest_index(chosen_y_mean, y_mean_bin)

            y = y_bin[index, :]
            t = temperatures[index]

            y_cont = baseline_als(y)
            y_cr = y / y_cont

            px = np.arange(n_px)

            #for plotting use only Loic's code
            # x_gd = gd21_waven(order, t, px)
            x_lt = lt21_waven(order, t, px)

            if PLOT_SPECTRA:
                # ax1.plot(x_gd, y_cr, color=colour, linestyle="--", label="%s GD" %hdf5_filename)
                ax1.plot(x_lt, y_cr, color=colour, linestyle="-", label="%s LT" %hdf5_filename)

            #do simulation only for first file
            if file_index == 0:
                nu_hr = gd21_waven(order, t, np.arange(-10.0, n_px + 10.0, 0.001))
                mol_hr = 1.0 - get_molecular_hr(molecule, nu_hr, Smin=smin)
                if PLOT_SPECTRA:
                    ax2.plot(nu_hr, mol_hr)

                abs_ix_hr = get_local_minima(mol_hr) #get hitran absorption minima indices
                abs_nu_hrs = nu_hr[abs_ix_hr] #get absorption nu
                abs_y_hrs = mol_hr[abs_ix_hr] #get absorption depth

                #N strongest lines
                abs_y_cutoff = sorted(abs_y_hrs)[n_lines] #select only the n strongest lines

                if PLOT_SPECTRA:
                    for abs_index, (abs_nu_hr, abs_y_hr) in enumerate(zip(abs_nu_hrs, abs_y_hrs)):
                        if abs_y_hr < abs_y_cutoff:
                            ax2.text(abs_nu_hr, abs_y_hr, abs_nu_hr)

            #loop through hr absorptions
            for abs_index, (abs_nu_hr, abs_y_hr) in enumerate(zip(abs_nu_hrs, abs_y_hrs)):
                if abs_y_hr < abs_y_cutoff:
                    # print("Checking line at %0.3f" %abs_nu_hr)

                    error = False

                    """use Loic's cal to find approx pixel range of absorption line"""
                    #find local minima close to the simulated lines
                    mol_gd_index = get_nearest_index(abs_nu_hr, x_lt)

                    #define micro window a few pixels to either side of expected pixel
                    local_indices_to_check = np.arange(mol_gd_index - 3, mol_gd_index + 4, 1)

                    #avoid lines at edge of detector
                    if min(local_indices_to_check) > 10 and max(local_indices_to_check) < n_px - 8:

                        #find local minima in pixels close to expected location of absorption line
                        local_minimum_index = get_local_minima(y_cr[local_indices_to_check]) + local_indices_to_check[0]

                        if len(local_minimum_index) == 1:
                            #if just one local minima found in spectrum micro window, centre indices on that
                            local_minimum_indices = np.arange(local_minimum_index - 1, local_minimum_index + 2, 1) #just take 1 point +- minimum

                            #find pixel number minimum
                            # x_hr, y_hr, px_min_position, chisq = fit_gaussian_absorption(px[local_minimum_indices], y_cr[local_minimum_indices], error=True)
                            x_hr, y_hr, px_min_position = fit_polynomial(px[local_minimum_indices], y_cr[local_minimum_indices], error=False)

                            if len(x_hr) > 0: #if no fitting error

                                # print("Single line found in data at %0.3f, chisq %0.3f" %(px_min_position, chisq*1.0e6))
                                abs_depth = np.min(y_hr)
                                abs_max = np.max(y_hr)

                                if abs_index == 0:
                                    label = "Gaussian fit to absorption bands"
                                else:
                                    label = ""

                                """Goddard"""
                                #convert pixel minimum and x_hr fit wavenumbers
                                nu_min_position_gd = gd21_waven(order, t, px_min_position)
                                delta_nu_gd = nu_min_position_gd - abs_nu_hr

                                """Loic"""
                                #convert pixel minimum and x_hr fit wavenumbers
                                nu_min_position_lt = lt21_waven(order, t, px_min_position)
                                delta_nu_lt = nu_min_position_lt - abs_nu_hr

                                if np.abs(delta_nu_gd) > 0.3 or np.abs(delta_nu_lt) > 0.3:
                                    print("Ignoring %0.1f line, too far from expected value" %abs_nu_hr)
                                    error = True

                                else:
                                    #plot
                                    if PLOT_SPECTRA:
                                        # x_hr_nu_gd = gd21_waven(order, t, x_hr)
                                        x_hr_nu_lt = lt21_waven(order, t, x_hr)
                                        # ax1.axvline(x=nu_min_position_gd, color=colour)
                                        # ax1.plot(x_hr_nu_gd, y_hr, color=colour, linestyle=":", label=label)
                                        ax1.axvline(x=nu_min_position_lt, color=colour)
                                        ax1.plot(x_hr_nu_lt, y_hr, color=colour, linestyle=":", label=label)

                            else:
                                # print("%s: fitting failed" %(hdf5_filename))
                                error = True
                        else:
                            # print("%s: %i lines found" %(hdf5_filename, len(local_minimum_index)))
                            error = True

                    else:
                        # print("%s: line too close to detector edge" %(hdf5_filename))
                        error = True

                    if not error:
                        text = "%s\t%i\t%0.3f\t%0.3f\t%0.3f\t%0.3f\t%0.3f\t%0.3f\t%0.3f\t%0.3f\t%0.3f" \
                            %(hdf5_filename, order, t, abs_nu_hr, px_min_position, nu_min_position_gd, nu_min_position_lt, \
                              delta_nu_gd, delta_nu_lt, \
                              abs_depth, abs_max)
                        write_log(text)

            hdf5_file.close()

            # abs_mean_gd = np.mean(np.abs(delta_nus_gd))
            # abs_mean_lt = np.mean(np.abs(delta_nus_lt))

            # # text = "%s\t%0.3f\t%0.3f" %(hdf5_filename, abs_mean_gd, abs_mean_lt)

            # print(text)
            # write_log(text)

        # ax1.legend()
        if PLOT_SPECTRA:
            plt.savefig("order_%i.png" %order)

# make_spectral_comparison_log(order_dict, chosen_bin, chosen_y_mean)

df = pd.read_csv(os.path.join(paths["BASE_DIRECTORY"], "reference_files", "so_spectral_cal_absorption_lines.txt"), sep="\t")

temperature_range = np.arange(np.floor(min(df["Temperature"])), np.ceil(max(df["Temperature"])) + temperature_grid_spacing, temperature_grid_spacing)
px_range = np.arange(320)

df["temp_bins"] = pd.cut(df["Temperature"], temperature_range, labels=False)
df["pixel_bins"] = pd.cut(df["Absorption pixel"], px_range, labels=False)

df["GD-LT"] = np.abs(df["GD delta nu"]) - np.abs(df["LT delta nu"])

with PdfPages("so_spectral_calibration_comparison.pdf") as pdf: #open pdf

    """plot all orders"""
    # grid = np.zeros((len(temperature_range), len(px_range))) + np.nan
    grid_list = [[[] for _ in range(len(px_range))] for _ in range(len(temperature_range))]

    for i, v in enumerate(df["GD-LT"]):
        # grid[df["temp_bins"][i], df["pixel_bins"][i]] = v
        # if 0.8 < df["Fitted max"][i] < 1.1:
        #     if not ((df["GD delta nu"][i] < 0.02) & (df["GD delta nu"][i] > -0.02)):
        #         if not ((df["LT delta nu"][i] < 0.02) & (df["LT delta nu"][i] > -0.02)):
                    grid_list[int(df["temp_bins"][i])][int(df["pixel_bins"][i])].append(v)

    grid_median = np.zeros((len(temperature_range), len(px_range))) + np.nan

    grid_counts = np.zeros((len(temperature_range), len(px_range)))
    for i in range(len(px_range)):
        for j in range(len(temperature_range)):
            if len(grid_list[j][i]) > 0:
                grid_median[j, i] = np.median(grid_list[j][i])
            grid_counts[j, i] = len(grid_list[j][i])

    plt.figure(figsize=(10, 5), constrained_layout=True)
    im = plt.pcolormesh(px_range, temperature_range, grid_median, cmap=redgreyblue, vmin=-0.1, vmax=0.1)
    plt.xlabel("Pixel Number")
    plt.ylabel("Interpolated Temperature")
    cbar = plt.colorbar(im)
    cbar.set_label("abs(delta GD)-abs(delta LT)\n(median if more than 1 value per grid point)", rotation=270, labelpad=20)
    plt.title("Comparing spectral calibration methods: Loic vs Goddard")

    pdf.savefig()
    plt.close()

    """plot each order"""
    for order in order_dict.keys():

        # grid = np.zeros((len(temperature_range), len(px_range))) + np.nan
        grid_list = [[[] for _ in range(len(px_range))] for _ in range(len(temperature_range))]

        for i in df.index[df["Order"] == order]:
            # grid[df["temp_bins"][i], df["pixel_bins"][i]] = v
            # if 0.8 < df["Fitted max"][i] < 1.1:
            #     if not ((df["GD delta nu"][i] < 0.02) & (df["GD delta nu"][i] > -0.02)):
            #         if not ((df["LT delta nu"][i] < 0.02) & (df["LT delta nu"][i] > -0.02)):
                        grid_list[int(df["temp_bins"][i])][int(df["pixel_bins"][i])].append(df["GD-LT"][i])

        grid_median = np.zeros((len(temperature_range), len(px_range))) + np.nan

        grid_counts = np.zeros((len(temperature_range), len(px_range)))
        for i in range(len(px_range)):
            for j in range(len(temperature_range)):
                if len(grid_list[j][i]) > 0:
                    grid_median[j, i] = np.median(grid_list[j][i])
                grid_counts[j, i] = len(grid_list[j][i])

        plt.figure(figsize=(10, 5), constrained_layout=True)
        im = plt.pcolormesh(px_range, temperature_range, grid_median, cmap=redgreyblue, vmin=-0.1, vmax=0.1)
        plt.xlabel("Pixel Number")
        plt.ylabel("Interpolated Temperature")
        cbar = plt.colorbar(im)
        cbar.set_label("abs(delta GD)-abs(delta LT)\n(median if more than 1 value per grid point)", rotation=270, labelpad=20)
        plt.title("Comparing spectral calibration methods: Loic vs Goddard Order %i" %order)

        pdf.savefig()
        plt.close()

# plt.figure(constrained_layout=True)
# im = plt.pcolormesh(px_range, temperature_range, grid_counts, cmap="binary")
# plt.xlabel("Pixel Number")
# plt.ylabel("Interpolated Temperature")
# cbar = plt.colorbar(im)
# cbar.set_label("Number of values per grid point", rotation=270, labelpad=20)
# plt.title("Comparing spectral calibration methods: Loic vs Goddard\nNumber of values per temperature-pixel bin")

# plt.figure(constrained_layout=True)
# plt.plot(np.sum(grid_counts, axis=1), temperature_range)
# plt.xlabel("Number of observations at this temperature")
# plt.ylabel("Interpolated Temperature")
# plt.title("Comparing spectral calibration methods: Loic vs Goddard\nNumber of values per temperature-pixel bin")

# from scipy.interpolate import griddata
# points = np.asfarray([df["Absorption pixel"], df["Temperature"]]).T
# values = df["GD-LT"]
# grid_x, grid_y = np.mgrid[0:320:1, np.floor(min(df["Temperature"])):np.ceil(max(df["Temperature"])):0.001]
# grid2 = griddata(points, values, (grid_x, grid_y), method='linear')

# plt.figure(constrained_layout=True)
# im2 = plt.imshow(grid2.T, aspect="auto", cmap="winter", \
#                  extent=(min(df["Absorption pixel"]), max(df["Absorption pixel"]), min(df["Temperature"]), max(df["Temperature"])), \
#                  vmin=-0.1, vmax=0.1, interpolation=None)
# cbar2 = plt.colorbar(im2)
# cbar2.set_label("abs(delta GD)-abs(delta LT)", rotation=270, labelpad=20)