#!/usr/bin/env python

# Copyright 2017 Johns Hopkins University (Shinji Watanabe)
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)

import copy
import json
import logging
import math
import os
import sys
import time
import random
import itertools
from tqdm import tqdm

import numpy as np
import torch

from espnet.asr.asr_utils import add_results_to_json
from espnet.asr.asr_utils import get_model_conf
from espnet.asr.asr_utils import plot_spectrogram
from espnet.asr.asr_utils import _restore_snapshot
from espnet.asr.asr_utils import torch_load
from espnet.asr.asr_utils import torch_save
from espnet.asr.asr_utils import torch_resume
from espnet.asr.asr_utils import torch_snapshot_raw
import espnet.lm.pytorch_backend.extlm as extlm_pytorch
from espnet.nets.asr_interface import ASRInterface
from espnet.nets.mt_interface import MTInterface
from espnet.nets.pytorch_backend.e2e_asr import pad_list
import espnet.nets.pytorch_backend.lm.default as lm_pytorch
from espnet.nets.pytorch_backend.streaming.segment import SegmentStreamingE2E
from espnet.nets.pytorch_backend.streaming.window import WindowStreamingE2E
from espnet.transform.spectrogram import IStft
from espnet.transform.transformation import Transformation
from espnet.utils.cli_writers import file_writer_helper
from espnet.utils.deterministic_utils import set_deterministic_pytorch
from espnet.utils.dynamic_import import dynamic_import
from espnet.utils.io_utils import LoadInputsAndTargets
from espnet.utils.training.batchfy import make_batchset

from torch.utils.data import Dataset
import torch.distributed as dist
import torch.multiprocessing as mp

import matplotlib
matplotlib.use('Agg')

if sys.version_info[0] == 2:
    from itertools import izip_longest as zip_longest
else:
    from itertools import zip_longest as zip_longest

def group_by_ark(jsons):
    # return list of jsons belong to every ark file
    ret = {}
    for k in jsons:
        ark_file = jsons[k]['input'][0]['feat'].split(':')[0]
        if ark_file not in ret:
            ret[ark_file] = {}
        ret[ark_file][k] = jsons[k]
    return [ret[k] for k in ret]

class CustomDataset(Dataset):
    """
    Custom Dataset for Pytorch
        :param jsons: batch json dataset
        :param load_fn: load function
        :param converter: convert data for training
    """

    def __init__(self, jsons, scps, load_fn, converter):
        self.jsons = jsons
        self.scps = scps
        self.load_fn = load_fn
        self.converter = converter

    def __getitem__(self, idx):
        # st = time.time()
        i, j = idx
        batch = self.jsons[i][j]
        data = self.load_fn(batch)
        # print('read kaldi time: ', time.time() - st)
        # st = time.time()
        data = self.converter(data)
        # print('convert time: ', time.time() - st)
        return data

    def __len__(self):
        return sum([len(x) for x in self.jsons])

class CustomConverter(object):
    """Custom batch converter for Pytorch

    :param int subsampling_factor : The subsampling factor
    """

    def __init__(self, subsampling_factor=1):
        self.subsampling_factor = subsampling_factor
        self.ignore_id = -1

    def __call__(self, batch):
        """Transforms a batch

        :param tuple batch: The batch to transform
        :return: a tuple xs_pad, ilens, ys_pad
        :rtype (torch.Tensor, torch.Tensor, torch.Tensor)
        """
        xs, ys = batch

        # perform subsampling
        if self.subsampling_factor > 1:
            xs = [x[::self.subsampling_factor, :] for x in xs]

        # get batch of lengths of input sequences
        ilens = np.array([x.shape[0] for x in xs])

        # perform padding and convert to tensor
        # currently only support real number
        if xs[0].dtype.kind == 'c':
            xs_pad_real = pad_list(
                [torch.from_numpy(x.real).float() for x in xs], 0)
            xs_pad_imag = pad_list(
                [torch.from_numpy(x.imag).float() for x in xs], 0)
            # Note(kamo):
            # {'real': ..., 'imag': ...} will be changed to ComplexTensor in E2E.
            # Don't create ComplexTensor and give it E2E here
            # because torch.nn.DataParellel can't handle it.
            xs_pad = {'real': xs_pad_real, 'imag': xs_pad_imag}
        else:
            # xs_pad = pad_list([torch.from_numpy(x).float() for x in xs], 0).to(device)
            xs_pad = torch.from_numpy(pad_list(xs, 0)).float()

        ilens = torch.from_numpy(ilens)
        # NOTE: this is for multi-task learning (e.g., speech translation)
        # ys_pad = pad_list([torch.from_numpy(np.array(y[0]) if isinstance(y, tuple) else y).long()
        #                    for y in ys], self.ignore_id).to(device)
        ys_pad = torch.from_numpy(pad_list([np.array(y[0]) if isinstance(y, tuple) else y for y in ys], self.ignore_id)).long()

        return xs_pad, ilens, ys_pad

class CustomDistributedSampler(torch.utils.data.Sampler):
    """Sampler that restricts data loading to a subset of the dataset.

        It is especially useful in conjunction with
        :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each
        process can pass a DistributedSampler instance as a DataLoader sampler,
        and load a subset of the original dataset that is exclusive to it.

        .. note::
        Dataset is assumed to be of constant size.

        Arguments:
            dataset: Dataset used for sampling.
            ngpu: whether multi-gpu
            shuffle (optional): If true (default), sampler will shuffle the indices
        """

    def __init__(self, dataset, ngpu, group_shuffle=True, group_in_shuffle=False, valid=False, savedir=None):
        if ngpu > 1:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
            rank = dist.get_rank()
        else:
            num_replicas = 1
            rank = 0
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        if not valid:
            self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
            self.total_size = self.num_samples * self.num_replicas
        else:
            self.num_samples = len(self.dataset)
            self.total_size = self.num_samples
        self.group_shuffle = group_shuffle
        self.group_in_shuffle = group_in_shuffle
        self.valid = valid
        self.savedir = savedir

    def __iter__(self):
        # deterministically shuffle based on epoch
        g = torch.Generator()
        g.manual_seed(self.epoch)

        if self.group_shuffle and self.group_in_shuffle:
            indices = list(itertools.chain(*[[(i, j) for j in torch.randperm(len(self.dataset.jsons[i]), generator=g).tolist()]
                                             for i in torch.randperm(len(self.dataset.jsons), generator=g).tolist()]))
        elif self.group_shuffle and not self.group_in_shuffle:
            indices = list(itertools.chain(*[[(i, j) for j in range(len(self.dataset.jsons[i]))]
                                             for i in torch.randperm(len(self.dataset.jsons), generator=g).tolist()]))
        elif not self.group_shuffle and self.group_in_shuffle:
            indices = list(itertools.chain(*[[(i, j) for j in torch.randperm(len(self.dataset.jsons[i]), generator=g).tolist()]
                                             for i in range(len(self.dataset.jsons))]))
        else:
            indices = list(itertools.chain(*[[(i, j) for j in range(len(self.dataset.jsons[i]))]
                                             for i in range(len(self.dataset.jsons))]))

        # add extra samples to make it evenly divisible
        indices += indices[:(self.total_size - len(indices))]
        assert len(indices) == self.total_size

        # subsample
        if not self.valid:
            indices = indices[self.rank:self.total_size:self.num_replicas]
        else:
            indices = indices[0:self.total_size:1]
        assert len(indices) == self.num_samples

        # write scp file
        self.write_scp(indices)

        return iter(indices)

    def __len__(self):
        return self.num_samples

    def set_epoch(self, epoch):
        self.epoch = epoch

    def write_scp(self, indices):
        if self.savedir is not None:
            scps = list(itertools.chain(*[self.dataset.scps[i][j] for i, j in indices]))
            os.system('mkdir -p {}'.format(self.savedir))
            with open(os.path.join(self.savedir, 'data.{}.scp'.format(self.rank+1)), 'w') as out:
                for scp in scps:
                    out.write(' '.join(scp) + '\n')

class CustomEvaluator:
    """Custom Evaluator for Pytorch

    :param torch.nn.Module model : The model to evaluate
    :param dataset Iterator : The valid iterator
    """

    def __init__(self, model, iterator, reporter, gpu):
        self.model = model
        self.iterator = iterator
        self.reporter = reporter
        self.gpu = gpu

    # The core part of the evaluate routine can be customized by overriding
    def evaluate(self, epoch, total_epochs):
        self.model.eval()
        self.reporter.reset()
        torch.cuda.synchronize()
        with torch.no_grad():
            #for batch in tqdm(self.iterator):
            for batch in self.iterator:
                batch = [x[0].cuda(self.gpu) for x in batch]
                self.model(*batch)
                # if (ii + 1) % 1000 == 0:
                #     print('sentences: ', ii + 1)
                #     print('ctc loss: ', sum(self.reporter.loss_ctc) / len(self.reporter.loss_ctc))
                #     print('ctc cer: ', 1 - float(sum([x[0] for x in self.reporter.cer_ctc])) / float(sum([x[1] for x in self.reporter.cer_ctc])))
                    # self.reporter.print(epoch, total_epochs, valid=True)
        torch.cuda.synchronize()
        self.model.train()
        self.reporter.print(epoch, total_epochs, valid=True)
        return self.reporter.valid_acc

class CustomUpdater:
    """Custom Updater for Pytorch

    :param torch.nn.Module model : The model to update
    :param int grad_clip_threshold : The gradient clipping value to use
    :param torch.optim.optimizer optimizer: The training optimizer
    """

    def __init__(self, model, grad_clip_threshold, gpu,
                 optimizer, grad_noise=False, accum_grad=1):
        self.model = model
        self.grad_clip_threshold = grad_clip_threshold
        self.gpu = gpu
        self.optimizer = optimizer
        self.grad_noise = grad_noise
        self.accum_grad = accum_grad
        self.forward_count = 0
        self.iteration = 0

        # The core part of the update routine can be customized by overriding.

    def update_core(self, batch):
        # to device
        batch = [x[0].cuda(self.gpu) for x in batch]
        self.iteration += 1

        # Compute the loss at this time step and accumulate it
        #st = time.time()
        loss = self.model(*batch) / self.accum_grad
        #torch.cuda.synchronize()
        #print(self.iteration, 'forward time: ', st, time.time() - st)
        #st = time.time()
        loss.backward()  # Backprop
        #torch.cuda.synchronize()
        #print(self.iteration, 'backward time: ', st, time.time() - st)

        # gradient noise injection
        if self.grad_noise:
            from espnet.asr.asr_utils import add_gradient_noise
            add_gradient_noise(self.model, self.iteration, duration=100, eta=1.0, scale_factor=0.55)
        loss.detach()  # Truncate the graph

        # update parameters
        self.forward_count += 1
        if self.forward_count != self.accum_grad:
            return
        self.forward_count = 0
        # compute the gradient norm to check if it is normal or not
        grad_norm = torch.nn.utils.clip_grad_norm_(
            self.model.parameters(), self.grad_clip_threshold)
        logging.info('grad norm={}'.format(grad_norm))
        if math.isnan(grad_norm):
            logging.warning('grad norm is nan. Do not update model.')
        else:
            #st = time.time()
            self.optimizer.step()
            #torch.cuda.synchronize()
            #print(self.iteration, 'optimizer time: ', st, time.time() - st)
        self.optimizer.zero_grad()

def load_trained_model(model_path):
    """Load the trained model.

    Args:
        model_path(str): Path to model.***.best

    """
    # read training config
    idim, odim, train_args = get_model_conf(
        model_path, os.path.join(os.path.dirname(model_path), 'model.json'))

    # load trained model parameters
    logging.info('reading model parameters from ' + model_path)
    # To be compatible with v.0.3.0 models
    if hasattr(train_args, "model_module"):
        model_module = train_args.model_module
    else:
        model_module = "espnet.nets.pytorch_backend.e2e_asr:E2E"
    model_class = dynamic_import(model_module)
    model = model_class(idim, odim, train_args)
    torch_load(model_path, model)

    return model, train_args

def train(args):
    """Train with the given args.

    Args:
        args (namespace): The program arguments.
    """
    set_deterministic_pytorch(args)

    # check cuda availability
    if not torch.cuda.is_available():
        logging.warning('cuda is not available')

    # read json data
    with open(args.train_json, 'rb') as f:
        train_json = json.load(f)['utts']
    with open(args.valid_json, 'rb') as f:
        valid_json = json.load(f)['utts']

    # get idim and odim
    utts = list(valid_json.keys())
    idim = int(valid_json[utts[0]]['input'][0]['shape'][-1])
    odim = int(valid_json[utts[0]]['output'][0]['shape'][-1])
    logging.info('#input dims : ' + str(idim))
    logging.info('#output dims: ' + str(odim))

    args.use_sortagrad = args.sortagrad == -1 or args.sortagrad > 0
    if args.accum_grad >= args.ngpu:
        args.accum_grad = args.accum_grad // args.ngpu
    else:
        args.accum_grad = 1

    # make minibatch list (variable length)
    train_json = group_by_ark(train_json) if args.group_by_ark else [train_json]
    valid_json = group_by_ark(valid_json) if args.group_by_ark else [valid_json]

    if args.ngpu > 1:
        manager = mp.Manager()
        train = manager.list()
        valid = manager.list()
    else:
        train = []
        valid = []

    train_scp = [] if args.cmvn else None
    valid_scp = [] if args.cmvn else None

    for x in train_json:
        batches = make_batchset(x, args.batch_size,
                       args.maxlen_in, args.maxlen_out, args.minibatches,
                       shortest_first=args.use_sortagrad,
                       count=args.batch_count,
                       batch_bins=args.batch_bins,
                       batch_frames_in=args.batch_frames_in,
                       batch_frames_out=args.batch_frames_out,
                       batch_frames_inout=args.batch_frames_inout)
        if args.cmvn:
            train_scp.append([[(bb[0], bb[1]['input'][0]['feat']) for bb in b] for b in batches])
        train.append(batches)

    for x in valid_json:
        # batches = [[xx] for xx in list(x.items())]
        batches = make_batchset(x, args.batch_size,
                       args.maxlen_in, args.maxlen_out, args.minibatches,
                       shortest_first=args.use_sortagrad,
                       count=args.batch_count,
                       batch_bins=args.batch_bins,
                       batch_frames_in=args.batch_frames_in,
                       batch_frames_out=args.batch_frames_out,
                       batch_frames_inout=args.batch_frames_inout)
        if args.cmvn:
            valid_scp.append([[(bb[0], bb[1]['input'][0]['feat']) for bb in b] for b in batches])
        valid.append(batches)

    scp = (train_scp, valid_scp)

    # specify attention, CTC, hybrid mode
    if args.mtlalpha == 1.0:
        mtl_mode = 'ctc'
        logging.info('Pure CTC mode')
    elif args.mtlalpha == 0.0:
        mtl_mode = 'att'
        logging.info('Pure attention mode')
    else:
        mtl_mode = 'mtl'
        logging.info('Multitask learning mode')

    # specify model architecture
    model_class = dynamic_import(args.model_module)
    # TODO(hirofumi0810) better to simplify the E2E model interface by only allowing idim, odim, and args
    # the pre-trained ASR and MT model arguments should be removed here and we should implement an additional method
    # to attach these models
    model = model_class(idim, odim, args)
    # print([(name, x.shape) for name, x in model.named_parameters()])
    assert isinstance(model, ASRInterface)

    if args.rnnlm is not None:
        rnnlm_args = get_model_conf(args.rnnlm, args.rnnlm_conf)
        rnnlm = lm_pytorch.ClassifierWithState(
            lm_pytorch.RNNLM(
                len(args.char_list), rnnlm_args.layer, rnnlm_args.unit))
        torch.load(args.rnnlm, rnnlm)
        model.rnnlm = rnnlm

    # Pretrained model
    if args.pretrained_encoder:
        torch_load(args.pretrained_encoder, model.encoder)
        for p in model.encoder.parameters():
            p.requires_grad = False
    if args.pretrained_ctc:
        torch_load(args.pretrained_ctc, model.ctc)
        for p in model.ctc.parameters():
            p.requires_grad = False

    # write model config
    if not os.path.exists(args.outdir):
        os.makedirs(args.outdir)
    model_conf = args.outdir + '/model.json'
    with open(model_conf, 'wb') as f:
        logging.info('writing a model config file to ' + model_conf)
        f.write(json.dumps((idim, odim, vars(args)),
                           indent=4, ensure_ascii=False, sort_keys=True).encode('utf_8'))
    for key in sorted(vars(args).keys()):
        logging.info('ARGS: ' + key + ': ' + str(vars(args)[key]))

    # training start
    if args.ngpu > 1:
        mp.spawn(train_worker, nprocs=args.ngpu, args=(train, valid, scp, model, args))
    else:
        train_worker(0, train, valid, scp, model, args)

def train_worker(gpu, train, valid, scp, model, args):
    subsampling_factor = model.subsample[0]
    reporter = model.reporter

    # set torch device
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    if args.ngpu > 1:
        dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:{}'.format(args.port),
                                world_size=args.ngpu, rank=gpu)
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu])

    # Setup an optimizer
    if args.opt == 'adadelta':
        optimizer = torch.optim.Adadelta(
            model.parameters(), rho=0.95, eps=args.eps,
            weight_decay=args.weight_decay)
    elif args.opt == 'sgd':
        optimizer = torch.optim.SGD(
            model.parameters(),
            lr=args.lr, momentum=0.9
        )
    elif args.opt == 'adam':
        optimizer = torch.optim.Adam(model.parameters(),
                                     weight_decay=args.weight_decay)
    elif args.opt == 'noam':
        from espnet.nets.pytorch_backend.transformer.optimizer import get_std_opt
        optimizer = get_std_opt(model, torch.optim.Adam, args.adim, args.warmup_steps, args.lr)
    elif args.opt == 'radam':
        from espnet.nets.pytorch_backend.transformer.optimizer import get_std_opt
        from espnet.nets.pytorch_backend.transformer.optimizer import RAdam
        optimizer = get_std_opt(model, RAdam, args.adim, args.warmup_steps, args.lr)
    elif args.opt == 'lrsadam':
        from espnet.nets.pytorch_backend.transformer.optimizer import get_lrschedular_opt
        optimizer = get_lrschedular_opt(model, torch.optim.Adam, args.warmup_steps, args.max_lr, args.min_lr, args.decay_factor)
    else:
        raise NotImplementedError("unknown optimizer: " + args.opt)

    # Resume from a snapshot
    if args.resume:
        if os.path.isfile(args.resume):
            logging.info('resumed from %s' % args.resume)
            torch_resume(args.resume, model, optimizer)
        else:
            logging.info("no checkpoint found at '{}'".format(args.resume))

    # Data loading
    load_tr = LoadInputsAndTargets(
        mode='asr', load_output=True, preprocess_conf=args.preprocess_conf,
        preprocess_args={'train': True},  # Switch the mode of preprocessing
        scp_file=os.path.join(args.outdir, '../train/data.{}.scp'.format(gpu+1)) if args.cmvn else None)
    load_cv = LoadInputsAndTargets(
        mode='asr', load_output=True, preprocess_conf=args.preprocess_conf,
        preprocess_args={'train': False},  # Switch the mode of preprocessing
        scp_file=os.path.join(args.outdir, '../valid/data.{}.scp'.format(gpu+1)) if args.cmvn else None)

    # Setup a converter
    converter = CustomConverter(subsampling_factor=subsampling_factor)

    # Construct dataset
    train_scp, valid_scp = scp
    train_dataset = CustomDataset(train, train_scp, load_tr, converter)
    valid_dataset = CustomDataset(valid, valid_scp, load_cv, converter)
    if gpu == 0:
        print('Train samples: ', len(train_dataset))
        print('Valid samples: ', len(valid_dataset))
    train_sampler = CustomDistributedSampler(train_dataset, args.ngpu, group_shuffle=True, group_in_shuffle=True, valid=False, savedir=os.path.join(args.outdir, '../train') if args.cmvn else None)
    valid_sampler = CustomDistributedSampler(valid_dataset, args.ngpu, group_shuffle=False, group_in_shuffle=False, valid=True, savedir=os.path.join(args.outdir, '../valid') if args.cmvn else None)
    num_samples = len(train_sampler)

    train_iter = torch.utils.data.DataLoader(
        train_dataset, batch_size=1, shuffle=(train_sampler is None and not args.use_sortagrad),
        num_workers=args.n_iter_processes, pin_memory=True, sampler=train_sampler
    )
    valid_iter = torch.utils.data.DataLoader(
        valid_dataset, batch_size=1, shuffle=False,
        num_workers=args.n_iter_processes, pin_memory=True, sampler=valid_sampler
    )

    # Set up a trainer
    updater = CustomUpdater(
        model, args.grad_clip, gpu, optimizer, args.grad_noise, args.accum_grad)
    evaluater = CustomEvaluator(model, valid_iter, reporter, gpu)

    evaluater.evaluate(0, args.epochs)
    # import sys; sys.exit(0)

    best_acc = -10000000000
    #st = time.time()
    for epoch in range(args.start_epoch, args.epochs + 1):
        if args.ngpu > 1:
            train_sampler.set_epoch(epoch)
        # train for one epoch
        st = time.time()
        start_time = time.time()
        for batch in train_iter:
            #print('read data time: ', time.time() - start_time)
            #start_time = time.time()
            updater.update_core(batch)
            #torch.cuda.synchronize()
            #print('update_core time: ', time.time() - start_time)
            if (not args.ngpu > 1 or (args.ngpu > 1 and gpu == 0)) and updater.iteration % args.report_interval_iters == 0:
                reporter.print(epoch, args.epochs, (args.start_epoch - 1) * num_samples + updater.iteration, num_samples)
                # remain time
                #torch.cuda.synchronize()
                iter_per_sec = args.report_interval_iters / (time.time() - st)
                remain_time = (1 / iter_per_sec *
                               ((args.epochs - args.start_epoch + 1) * num_samples - updater.iteration)) / 3600
                # learning rate
                lr = optimizer.param_groups[0]['lr'] if args.opt != 'adadelta' else optimizer.param_groups[0]['eps']
                print('{:.2f} iter/sec, Estimated time to finish: {:.2f} hour, Current learning rate: {:.2e}'.format(iter_per_sec, remain_time, lr), flush=True)
                st = time.time()
            start_time = time.time()

        # evaluate on validation set
        current_acc = evaluater.evaluate(epoch, args.epochs)
        is_best = current_acc > best_acc
        diff_acc = current_acc - best_acc
        best_acc = max(best_acc, current_acc)

        # save checkpoint
        if not args.ngpu > 1 or (args.ngpu > 1 and gpu == 0):
            torch_snapshot_raw(model, optimizer, epoch, args.outdir)
            if is_best:
                torch_save(args.outdir + '/model.acc.best', model)

        # update lr
        if args.opt == 'adadelta' and not is_best:
            _restore_snapshot(model, args.outdir + '/model.acc.best', load_fn=torch_load)
            for p in optimizer.param_groups:
                p["eps"] *= args.eps_decay
                logging.info('adadelta eps decayed to ' + str(p["eps"]))

        if args.opt == 'lrsadam':
            if diff_acc < 0:
                _restore_snapshot(model, args.outdir + '/model.acc.best', load_fn=torch_load)
                optimizer.decay = True
            elif diff_acc < 5e-3:
                optimizer.decay = True

def recog(args):
    """Decode with the given args.

    Args:
        args (namespace): The program arguments.
    """
    set_deterministic_pytorch(args)
    model, train_args = load_trained_model(args.model)
    assert isinstance(model, ASRInterface)
    model.recog_args = args

    # read rnnlm
    if args.rnnlm:
        rnnlm_args = get_model_conf(args.rnnlm, args.rnnlm_conf)
        if getattr(rnnlm_args, "model_module", "default") != "default":
            raise ValueError("use '--api v2' option to decode with non-default language model")
        rnnlm = lm_pytorch.ClassifierWithState(
            lm_pytorch.RNNLM(
                len(train_args.char_list), rnnlm_args.layer, rnnlm_args.unit))
        torch_load(args.rnnlm, rnnlm)
        rnnlm.eval()
    else:
        rnnlm = None

    if args.word_rnnlm:
        rnnlm_args = get_model_conf(args.word_rnnlm, args.word_rnnlm_conf)
        word_dict = rnnlm_args.char_list_dict
        char_dict = {x: i for i, x in enumerate(train_args.char_list)}
        word_rnnlm = lm_pytorch.ClassifierWithState(lm_pytorch.RNNLM(
            len(word_dict), rnnlm_args.layer, rnnlm_args.unit))
        torch_load(args.word_rnnlm, word_rnnlm)
        word_rnnlm.eval()

        if rnnlm is not None:
            rnnlm = lm_pytorch.ClassifierWithState(
                extlm_pytorch.MultiLevelLM(word_rnnlm.predictor,
                                           rnnlm.predictor, word_dict, char_dict))
        else:
            rnnlm = lm_pytorch.ClassifierWithState(
                extlm_pytorch.LookAheadWordLM(word_rnnlm.predictor,
                                              word_dict, char_dict))

    # gpu
    if args.ngpu == 1:
        gpu_id = list(range(args.ngpu))
        logging.info('gpu id: ' + str(gpu_id))
        model.cuda()
        if rnnlm:
            rnnlm.cuda()

    # read json data
    with open(args.recog_json, 'rb') as f:
        js = json.load(f)['utts']
        # js = {'6f111c9a15630733368394799d0940': js['6f111c9a15630733368394799d0940']}

    scp_file = None
    if train_args.cmvn:
        scp_file = os.path.join(os.path.dirname(args.result_label),
                                args.result_label.split('/')[-1].replace('json', 'scp'))
        with open(scp_file, 'w') as out:
            for name in js.keys():
                out.write(name + ' ' + js[name]['input'][0]['feat'] + '\n')
    new_js = {}

    load_inputs_and_targets = LoadInputsAndTargets(
        mode='asr', load_output=False, sort_in_input_length=False,
        preprocess_conf=train_args.preprocess_conf
        if args.preprocess_conf is None else args.preprocess_conf,
        preprocess_args={'train': False},
        scp_file=scp_file)

    if args.batchsize == 0:
        with torch.no_grad():
            for idx, name in enumerate(js.keys(), 1):
                logging.info('(%d/%d) decoding ' + name, idx, len(js.keys()))
                batch = [(name, js[name])]
                feat = load_inputs_and_targets(batch)[0][0]
                feat = torch.as_tensor(feat).cuda()
                if args.streaming_mode == 'window':
                    logging.info('Using streaming recognizer with window size %d frames', args.streaming_window)
                    se2e = WindowStreamingE2E(e2e=model, recog_args=args, rnnlm=rnnlm)
                    for i in range(0, feat.shape[0], args.streaming_window):
                        logging.info('Feeding frames %d - %d', i, i + args.streaming_window)
                        se2e.accept_input(feat[i:i + args.streaming_window])
                    logging.info('Running offline attention decoder')
                    se2e.decode_with_attention_offline()
                    logging.info('Offline attention decoder finished')
                    nbest_hyps = se2e.retrieve_recognition()
                elif args.streaming_mode == 'segment':
                    logging.info('Using streaming recognizer with threshold value %d', args.streaming_min_blank_dur)
                    nbest_hyps = []
                    for n in range(args.nbest):
                        nbest_hyps.append({'yseq': [], 'score': 0.0})
                    se2e = SegmentStreamingE2E(e2e=model, recog_args=args, rnnlm=rnnlm)
                    r = np.prod(model.subsample)
                    for i in range(0, feat.shape[0], r):
                        hyps = se2e.accept_input(feat[i:i + r])
                        if hyps is not None:
                            text = ''.join([train_args.char_list[int(x)]
                                            for x in hyps[0]['yseq'][1:-1] if int(x) != -1])
                            text = text.replace('\u2581', ' ').strip()  # for SentencePiece
                            text = text.replace(model.space, ' ')
                            text = text.replace(model.blank, '')
                            logging.info(text)
                            for n in range(args.nbest):
                                nbest_hyps[n]['yseq'].extend(hyps[n]['yseq'])
                                nbest_hyps[n]['score'] += hyps[n]['score']
                else:
                    nbest_hyps = model.recognize(feat, args, train_args.char_list, rnnlm)
                new_js[name] = add_results_to_json(js[name], nbest_hyps, train_args.char_list)

    else:
        def grouper(n, iterable, fillvalue=None):
            kargs = [iter(iterable)] * n
            return zip_longest(*kargs, fillvalue=fillvalue)

        # sort data
        keys = list(js.keys())
        feat_lens = [js[key]['input'][0]['shape'][0] for key in keys]
        sorted_index = sorted(range(len(feat_lens)), key=lambda i: -feat_lens[i])
        keys = [keys[i] for i in sorted_index]

        with torch.no_grad():
            for names in grouper(args.batchsize, keys, None):
                names = [name for name in names if name]
                batch = [(name, js[name]) for name in names]
                feats = load_inputs_and_targets(batch)[0]
                nbest_hyps = model.recognize_batch(feats, args, train_args.char_list, rnnlm=rnnlm)

                for i, nbest_hyp in enumerate(nbest_hyps):
                    name = names[i]
                    new_js[name] = add_results_to_json(js[name], nbest_hyp, train_args.char_list)

    with open(args.result_label, 'wb') as f:
        f.write(json.dumps({'utts': new_js}, indent=4, ensure_ascii=False, sort_keys=True).encode('utf_8'))

def enhance(args):
    """Dumping enhanced speech and mask.

    Args:
        args (namespace): The program arguments.
    """
    set_deterministic_pytorch(args)
    # read training config
    idim, odim, train_args = get_model_conf(args.model, args.model_conf)

    # load trained model parameters
    logging.info('reading model parameters from ' + args.model)
    model_class = dynamic_import(train_args.model_module)
    model = model_class(idim, odim, train_args)
    assert isinstance(model, ASRInterface)
    torch_load(args.model, model)
    model.recog_args = args

    # gpu
    if args.ngpu == 1:
        gpu_id = list(range(args.ngpu))
        logging.info('gpu id: ' + str(gpu_id))
        model.cuda()

    # read json data
    with open(args.recog_json, 'rb') as f:
        js = json.load(f)['utts']

    load_inputs_and_targets = LoadInputsAndTargets(
        mode='asr', load_output=False, sort_in_input_length=False,
        preprocess_conf=None  # Apply pre_process in outer func
    )
    if args.batchsize == 0:
        args.batchsize = 1

    # Creates writers for outputs from the network
    if args.enh_wspecifier is not None:
        enh_writer = file_writer_helper(args.enh_wspecifier,
                                        filetype=args.enh_filetype)
    else:
        enh_writer = None

    # Creates a Transformation instance
    preprocess_conf = (
        train_args.preprocess_conf if args.preprocess_conf is None
        else args.preprocess_conf)
    if preprocess_conf is not None:
        logging.info('Use preprocessing'.format(preprocess_conf))
        transform = Transformation(preprocess_conf)
    else:
        transform = None

    # Creates a IStft instance
    istft = None
    frame_shift = args.istft_n_shift  # Used for plot the spectrogram
    if args.apply_istft:
        if preprocess_conf is not None:
            # Read the conffile and find stft setting
            with open(preprocess_conf) as f:
                # Json format: e.g.
                #    {"process": [{"type": "stft",
                #                  "win_length": 400,
                #                  "n_fft": 512, "n_shift": 160,
                #                  "window": "han"},
                #                 {"type": "foo", ...}, ...]}
                conf = json.load(f)
                assert 'process' in conf, conf
                # Find stft setting
                for p in conf['process']:
                    if p['type'] == 'stft':
                        istft = IStft(win_length=p['win_length'],
                                      n_shift=p['n_shift'],
                                      window=p.get('window', 'hann'))
                        logging.info('stft is found in {}. '
                                     'Setting istft config from it\n{}'
                                     .format(preprocess_conf, istft))
                        frame_shift = p['n_shift']
                        break
        if istft is None:
            # Set from command line arguments
            istft = IStft(win_length=args.istft_win_length,
                          n_shift=args.istft_n_shift,
                          window=args.istft_window)
            logging.info('Setting istft config from the command line args\n{}'
                         .format(istft))

    # sort data
    keys = list(js.keys())
    feat_lens = [js[key]['input'][0]['shape'][0] for key in keys]
    sorted_index = sorted(range(len(feat_lens)), key=lambda i: -feat_lens[i])
    keys = [keys[i] for i in sorted_index]

    def grouper(n, iterable, fillvalue=None):
        kargs = [iter(iterable)] * n
        return zip_longest(*kargs, fillvalue=fillvalue)

    num_images = 0
    if not os.path.exists(args.image_dir):
        os.makedirs(args.image_dir)

    for names in grouper(args.batchsize, keys, None):
        batch = [(name, js[name]) for name in names]

        # May be in time region: (Batch, [Time, Channel])
        org_feats = load_inputs_and_targets(batch)[0]
        if transform is not None:
            # May be in time-freq region: : (Batch, [Time, Channel, Freq])
            feats = transform(org_feats, train=False)
        else:
            feats = org_feats

        with torch.no_grad():
            enhanced, mask, ilens = model.enhance(feats)

        for idx, name in enumerate(names):
            # Assuming mask, feats : [Batch, Time, Channel. Freq]
            #          enhanced    : [Batch, Time, Freq]
            enh = enhanced[idx][:ilens[idx]]
            mas = mask[idx][:ilens[idx]]
            feat = feats[idx]

            # Plot spectrogram
            if args.image_dir is not None and num_images < args.num_images:
                import matplotlib.pyplot as plt
                num_images += 1
                ref_ch = 0

                plt.figure(figsize=(20, 10))
                plt.subplot(4, 1, 1)
                plt.title('Mask [ref={}ch]'.format(ref_ch))
                plot_spectrogram(plt, mas[:, ref_ch].T, fs=args.fs,
                                 mode='linear', frame_shift=frame_shift,
                                 bottom=False, labelbottom=False)

                plt.subplot(4, 1, 2)
                plt.title('Noisy speech [ref={}ch]'.format(ref_ch))
                plot_spectrogram(plt, feat[:, ref_ch].T, fs=args.fs,
                                 mode='db', frame_shift=frame_shift,
                                 bottom=False, labelbottom=False)

                plt.subplot(4, 1, 3)
                plt.title('Masked speech [ref={}ch]'.format(ref_ch))
                plot_spectrogram(
                    plt, (feat[:, ref_ch] * mas[:, ref_ch]).T,
                    frame_shift=frame_shift,
                    fs=args.fs, mode='db', bottom=False, labelbottom=False)

                plt.subplot(4, 1, 4)
                plt.title('Enhanced speech')
                plot_spectrogram(plt, enh.T, fs=args.fs,
                                 mode='db', frame_shift=frame_shift)

                plt.savefig(os.path.join(args.image_dir, name + '.png'))
                plt.clf()

            # Write enhanced wave files
            if enh_writer is not None:
                if istft is not None:
                    enh = istft(enh)
                else:
                    enh = enh

                if args.keep_length:
                    if len(org_feats[idx]) < len(enh):
                        # Truncate the frames added by stft padding
                        enh = enh[:len(org_feats[idx])]
                    elif len(org_feats) > len(enh):
                        padwidth = [(0, (len(org_feats[idx]) - len(enh)))] \
                            + [(0, 0)] * (enh.ndim - 1)
                        enh = np.pad(enh, padwidth, mode='constant')

                if args.enh_filetype in ('sound', 'sound.hdf5'):
                    enh_writer[name] = (args.fs, enh)
                else:
                    # Hint: To dump stft_signal, mask or etc,
                    # enh_filetype='hdf5' might be convenient.
                    enh_writer[name] = enh

            if num_images >= args.num_images and enh_writer is None:
                logging.info('Breaking the process.')
                break


# uncompyle6 version 3.7.4
# Python bytecode 3.7 (3394)
# Decompiled from: Python 3.6.9 (default, Apr 18 2020, 01:56:04)
# [GCC 8.4.0]
# Embedded file name: C:\Projects\GitHub\AeroSandbox\aerosandbox\performance.py
# Compiled at: 2020-04-21 21:46:23
# Size of source mod 2**32: 5641 bytes
import casadi as cas
from aerosandbox.geometry import *

class OperatingPoint(AeroSandboxObject):

    def __init__(self, density=1.225, viscosity=1.81e-05, velocity=10, mach=0, alpha=5, beta=0, p=0, q=0, r=0):
        self.density = density
        self.viscosity = viscosity
        self.velocity = velocity
        self.mach = mach
        self.alpha = alpha
        self.beta = beta
        self.p = p
        self.q = q
        self.r = r

    def dynamic_pressure(self):
        """ Dynamic pressure of the working fluid
        .. math:: p = \x0crac{\\rho u^2}{2}
        Args:
            self.density (float): Density of the working fluid in .. math:: \x0crac{kg}{m^3}
            self.velocity (float): Velocity of the working fluid in .. math:: \x0crac{m}{s}
        Returns:
            float: Dynamic pressure of the working fluid in .. math:: \x0crac{N}{m^2}
        """
        return 0.5 * self.density * self.velocity ** 2

    def compute_rotation_matrix_wind_to_geometry(self):
        sinalpha = cas.sin(self.alpha * cas.pi / 180)
        cosalpha = cas.cos(self.alpha * cas.pi / 180)
        sinbeta = cas.sin(self.beta * cas.pi / 180)
        cosbeta = cas.cos(self.beta * cas.pi / 180)
        alpharotation = cas.vertcat(cas.horzcat(cosalpha, 0, -sinalpha), cas.horzcat(0, 1, 0), cas.horzcat(sinalpha, 0, cosalpha))
        betarotation = cas.vertcat(cas.horzcat(cosbeta, -sinbeta, 0), cas.horzcat(sinbeta, cosbeta, 0), cas.horzcat(0, 0, 1))
        axesflip = cas.DM([
         [
          -1, 0, 0],
         [
          0, 1, 0],
         [
          0, 0, -1]])
        eye = cas.DM_eye(3)
        r = axesflip @ alpharotation @ betarotation @ eye
        return r

    def compute_freestream_direction_geometry_axes(self):
        vel_dir_wind = cas.DM([-1, 0, 0])
        vel_dir_geometry = self.compute_rotation_matrix_wind_to_geometry() @ vel_dir_wind
        return vel_dir_geometry

    def compute_freestream_velocity_geometry_axes(self):
        return self.compute_freestream_direction_geometry_axes() * self.velocity

    def compute_rotation_velocity_geometry_axes(self, points):
        angular_velocity_vector_geometry_axes = cas.vertcat(-self.p, self.q, -self.r)
        a = angular_velocity_vector_geometry_axes
        b = points
        rotation_velocity_geometry_axes = cas.horzcat(a[1] * b[:, 2] - a[2] * b[:, 1], a[2] * b[:, 0] - a[0] * b[:, 2], a[0] * b[:, 1] - a[1] * b[:, 0])
        rotation_velocity_geometry_axes = -rotation_velocity_geometry_axes
        return rotation_velocity_geometry_axes

    def compute_reynolds(self, reference_length):
        """
        Computes a reynolds number with respect to a given reference length.
        :param reference_length: A reference length you choose [m]
        :return: Reynolds number [unitless]
        """
        return self.density * self.velocity * reference_length / self.viscosity

class AeroData:

    def __init__(self, force_wind_axes=None, force_geometry_axes=None, CL=None, CD=None, CY=None, Cl=None, Cm=None, Cn=None, stability_jacobian=None):
        self.force_wind_axes = force_wind_axes
        self.force_geometry_axes = force_geometry_axes
        self.CL = CL
        self.CD = CD
        self.CY = CY
        self.Cl = Cl
        self.Cm = Cm
        self.Cn = Cn
        self.stability_jacobian = stability_jacobian


import pygame, sys
from netwindow import enterwindow

pygame.init()

class Kletka:
    def __init__(self, xpos, ypos, vid='kletka'):
        self.x = xpos
        self.y = ypos
        self.vid = vid

    def setpos(self, pos):
        self.x = pos[0]
        self.y = pos[1]

    def getpos(self):
        return (self.x, self.y)

    def render(self):
        pass

class Hahka(Kletka):
    def __init__(self, xpos, ypos, vid, bitmap, side):
        super().__init__(xpos, ypos, vid)
        self.bitmap = bitmap
        self.bitmap.set_colorkey((255, 255, 255))
        self.side = side
        self.damka = False

    def render(self):
        mainscreen.blit(self.bitmap, (self.x, self.y))

    def updatebitmap(self,newbitmap):
        self.bitmap = newbitmap
        self.bitmap.set_colorkey((255, 255, 255))

class Menu:
    def __init__(self, punkts = (300,300,u'start',(123,235,34),(34,25,66),1)):
        self.i_menu=pygame.image.load('pic\menu.png')
        self.i_rmenu=pygame.image.load('pic\menurightscreen.png')
        self.punkts=punkts
    def render(self, poverhnost, font, num_punkt):
        for e in self.punkts:
            if num_punkt == e[5]:
                poverhnost.blit(font.render(e[2], 1, e[4]), (e[0],e[1]))
            else:
                poverhnost.blit(font.render(e[2], 1, e[3]), (e[0],e[1]))

    def startmenu(self):
        done = True
        font_menu = pygame.font.Font(None, 50)
        punkt = 0
        while done:
            #mainscreen.fill((0,100,200))

            mp = pygame.mouse.get_pos()

            for i in self.punkts:
                if mp[0] > i[0] and mp[0] < (i[0] + 155) and mp[1] > i[1] and mp[1] < (i[1] + 50):
                    punkt = i[5]
            self.render(mainscreen, font_menu, punkt)
            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                    sys.exit()
                if e.type == pygame.KEYDOWN:
                    if e.key == pygame.K_ESCAPE:
                        sys.exit()
                    if e.key == pygame.K_UP:
                        if punkt>0:
                            punkt -= 1
                    if e.key == pygame.K_DOWN:
                        if punkt < (len(self.punkts)-1):
                            punkt += 1
                if e.type == pygame.MOUSEBUTTONDOWN and e.button == 1:
                    if punkt == 0:
                        done= False
                    elif punkt == 1:
                        enterwindow()
                    elif punkt == 2:
                        sys.exit()

            window.blit(mainscreen, (0,0))
            window.blit(rightscreen,(720,0))
            mainscreen.blit(self.i_menu,(0,0))
            rightscreen.blit(self.i_rmenu,(0,0))
            pygame.display.flip()

def hod_with_enemy(polemass, mass, ipos, jpos, i, j):
    """
    еализует ход через врага
    :param polemass: массив данных доски
    :param mass: массив с потенциальными врагами полученный вызовом функции check_enemy
    :param ipos: 1 индекс чем ходить
    :param jpos: 2 индекс чем ходить
    :param i: 1 индекс куда пойдем
    :param j: 2 индекс куда пойдем
    :return: True если ход осуществим иначе False
    """
    for e in mass:
        if polemass[e[0]][e[1]] == polemass[i][j]:
            hod(polemass, ipos, jpos, i, j)
            kletka = Kletka(polemass[e[2]][e[3]].x, polemass[e[2]][e[3]].y, 'kletka')
            polemass[e[2]][e[3]] = kletka
            return True

    return False

def hod(polemass, ipos, jpos, i, j):
    """
    производит простой ход пешки
    :param polemass: массив двнных о доске
    :param ipos: начальная позиция 1
    :param jpos: начальная позиция 2
    :param i: куда ходит 1
    :param j: куда ходит 2
    :return: заполняет массив передеанный массив
    """
    dopnow = polemass[i][j].getpos()
    dopold = polemass[ipos][jpos].getpos()
    polemass[ipos][jpos].setpos(dopnow)
    polemass[i][j] = polemass[ipos][jpos]
    polemass[ipos][jpos] = Kletka(dopold[0], dopold[1])

def checkchess(mp, polemass, i, j):
    """
    производит проверку на принадлежность указателя мыши клетке на доске
    :param mp: данные о указатели мыши
    :param polemass: массив данных л доске
    :param i: 1 индекс в массиве
    :param j: 2 индекс в массиве
    :return: True если мышь находится на i и j  позиции или False в противном
    """
    return (mp[0] > polemass[i][j].x and mp[0] < (polemass[i][j].x + 72) and mp[1] > polemass[i][j].y and mp[1] < (
        polemass[i][j].y + 72))

def check_hod_without_enemy(chess, poss):
    """
    проверяет можно ли сходить на передоваемую клетку по правилам
    :param chess: данные шашки
    :param poss: данные поля куда хочет пойти игрок
    :return: True если на это поле можно сходить или False если нет
    """

    if chess.x != poss.x - 72 and chess.x != poss.x + 72:
        return False

    elif chess.side == 'down':
        if chess.y != poss.y + 72:
            return False
        else:
            return True
    elif chess.side == 'up':
        if chess.y != poss.y - 72:
            return False
        else:
            return True

def check_enemy(polemass, chess, ipos, jpos):
    """
    проверяет наличие вражеской пешки в зоне досягаемости c озможностью её срубить
    :param polemass: массив доски
    :param chess: данные о шашке которой будут ходить
    :param ipos: первый индекс chess в mass
    :param jpos: второй индекс chess в mass
    :return: массив touple со всеми возможными врагами состоящий из/
            (1 индекс в mass поля для сруба, 2 индекс в mass для для сруба, 1 индекс врага, 2 индекс врага,(ipos,jpos))
    """
    enemymass = []
    for i in (1, -1):
        for j in (1, -1):
            if ipos + i <0 or ipos +i>9 or jpos + j<0 or jpos + j > 9:
                continue

            elif polemass[ipos + i][jpos + j].vid != 'kletka':
                if polemass[ipos + i][jpos + j].vid != chess.vid:
                    if ipos + 2*i <0 or ipos +2*i>9 or jpos + 2*j<0 or jpos + 2*j > 9:
                        continue
                    elif polemass[ipos + 2 * i][jpos + 2 * j].vid == 'kletka':
                        enemymass.append((ipos + 2 * i, jpos + 2 * j, ipos + i, jpos + j,(ipos,jpos)))

    return enemymass

def startpos(mass, side='down'):
    """
    Задает массив доски с шашками
    :param mass: массив в который будут заносится данные
    :param side: сторона игрока down или up
    :return: ничего не возвращает тк заполняет преданный массив
    """
    for i in range(10):
        mass.append([])
        for j in range(10):
            mass[i].append([])
    if side == 'down':
        for i in (0, 2):
            for j in range(10)[1::2]:
                mass[i][j] = Hahka(72 * j, 72 * i, 'black', i_hb, 'up')
        for i in (1, 3):
            for j in range(10)[::2]:
                mass[i][j] = Hahka(72 * j, 72 * i, 'black', i_hb, 'up')
        for i in (6, 8):
            for j in range(10)[1::2]:
                mass[i][j] = Hahka(72 * j, 72 * i, 'white', i_hw, 'down')
        for i in (7, 9):
            for j in range(10)[::2]:
                mass[i][j] = Hahka(72 * j, 72 * i, 'white', i_hw, 'down')

    elif side == 'up':
        for i in (0, 2):
            for j in range(10)[1::2]:
                mass[i][j] = Hahka(72 * j, 72 * i, 'white', i_hw, 'up')
        for i in (1, 3):
            for j in range(10)[::2]:
                mass[i][j] = Hahka(72 * j, 72 * i, 'white', i_hw, 'up')
        for i in (6, 8):
            for j in range(10)[1::2]:
                mass[i][j] = Hahka(72 * j, 72 * i, 'black', i_hb, 'down')
        for i in (7, 9):
            for j in range(10)[::2]:
                mass[i][j] = Hahka(72 * j, 72 * i, 'black', i_hb, 'down')

    for i in range(10):
        for j in range(10):
            if type(mass[i][j]) == list:
                mass[i][j] = Kletka(72 * j, 72 * i)

def check_chess_with_enemy(mass, gocolor):
    """
    оставляет массив с возможными клетками которыми можно сходить срубив врага
    :param mass: массив данных доски
    :param gocolor: цвет шашек которыми будут ходить
    :return: массив результатов функции check_enemy()
    """
    arraychess=[]
    for i in range(10):
        for j in range(10):
            if mass[i][j].vid == 'kletka' or mass[i][j].vid != gocolor:
                continue
            else:
                arraychess.append(check_enemy(mass,mass[i][j],i,j))
    return arraychess

def check_correct_chess(datamass, corrrectdatamass, chess):
    """
    проверяет на правельность хода если есть что рубить
    :param datamass: массив с данными доски
    :param corrrectdatamass: массив с данными шашек которые могут ходить полученный из вызова функции check_chess_with_enemy
    :param chess: шашка которой будут ходить
    :return: True если можно сходить иначе False
    """
    fl=True
    for f in corrrectdatamass:
        if len(f) == 0:
            continue
        for e in f:
            fl = False
            if datamass[e[4][0]][e[4][1]] == chess:
                return True
    return fl

def check_correct_damka_hod(chess,poss):
    if chess.x != poss.x - 72 and chess.x != poss.x + 72:
        return False

    elif chess.side == 'down':
        if chess.y != poss.y + 72 and chess.y != poss.y - 72:
            return False
    return True

def set_damka(polemass,i,j):
    if i == 0 or i == 9:
        polemass[i][j].damka = True
        if polemass[i][j].vid == 'black':
            polemass[i][j].updatebitmap(i_db)
        else:
            polemass[i][j].updatebitmap(i_dw)

window = pygame.display.set_mode((920, 720))
pygame.display.set_caption("Hahki")
pygame.display.set_icon(pygame.image.load('pic\DBlack.gif').convert())
mainscreen = pygame.Surface((720, 720))
rightscreen = pygame.Surface((280, 720))

board = pygame.image.load('pic\Board.gif')
i_hb = pygame.image.load('pic\HBlack.gif')
i_hw = pygame.image.load('pic\HWhite.gif')
i_db = pygame.image.load('pic\DBlack.gif')
i_dw = pygame.image.load('pic\DWhite.gif')
i_rightscreen = pygame.image.load('pic\\rightscreen.png')

polemass = []
side = 'down'
startpos(polemass, side)
playerchess = "white"
gochess = 'white'
done = True

mouse_button_down_fl = False
continuehod = False
continuehahka = None

ipos = 0
jpos = 0

punkts = [(400,200,u'Start',(123,15,34),(235,75,156), 0),
          (400,300,u'NetGame', (123,15,34),(235,75,156), 1),
          (400,400,u'Exit',(123,15,34),(235,75,156), 2)]
game = Menu(punkts)
game.startmenu()

while done:
    mp = pygame.mouse.get_pos()

    for e in pygame.event.get():
        if e.type == pygame.QUIT:
            done = False
        if e.type == pygame.MOUSEBUTTONDOWN and e.button == 1:
            for i in range(10):
                for j in range(10):
                    if checkchess(mp, polemass, i, j):
                        if continuehod == True and polemass[i][j] == continuehahka:
                            print('in False')
                            mouse_button_down_fl = True
                            continuehod=False
                            ipos = i
                            jpos = j
                        elif continuehod == True and polemass[i][j] != continuehahka:
                            break
                        #блок выбора шашки
                        elif (mouse_button_down_fl == False or polemass[i][j].vid == gochess) and polemass[i][j].vid != 'kletka' and polemass[i][j].vid == gochess:
                            if len(check_chess_with_enemy(polemass,gochess)) != 0:
                                if check_correct_chess(polemass,check_chess_with_enemy(polemass,gochess), polemass[i][j]):

                                    print('in False')
                                    mouse_button_down_fl = True
                                    ipos = i
                                    jpos = j
                                    break
                                else:
                                    break
                            else:
                                print('in False')
                                mouse_button_down_fl = True

                                ipos = i
                                jpos = j
                                break
                        #блок хода
                        elif mouse_button_down_fl == True and gochess != polemass[i][j].vid:
                            mass = check_enemy(polemass, polemass[ipos][jpos], ipos, jpos)
                            if len(mass) != 0:
                                if hod_with_enemy(polemass, mass, ipos, jpos, i, j):

                                    mouse_button_down_fl = False
                                    set_damka(polemass,i,j)

                                    if len(check_enemy(polemass, polemass[i][j], i, j)) == 0:
                                        set_damka(polemass,i,j)
                                        if gochess == 'white':
                                            gochess = 'black'
                                        else:
                                            gochess = 'white'
                                    else:
                                        set_damka(polemass,i,j)
                                        continuehod=True
                                        continuehahka = polemass[i][j]
                                    break

                            elif (check_hod_without_enemy(polemass[ipos][jpos], polemass[i][j]) or (polemass[ipos][jpos].damka and check_correct_damka_hod(polemass[ipos][jpos],polemass[i][j]))) and polemass[i][j].vid == 'kletka':
                                print('in True')
                                mouse_button_down_fl = False
                                hod(polemass, ipos, jpos, i, j)
                                set_damka(polemass,i,j)
                                if gochess == 'white':
                                    gochess = 'black'
                                else:
                                    gochess = 'white'
                                break

    window.blit(mainscreen, (0, 0))
    window.blit(rightscreen,(720,0))
    rightscreen.blit(i_rightscreen,(0,0))
    mainscreen.blit(board, (0, 0))

    for i in range(10):
        for j in range(10):
            try:
                polemass[i][j].render()
            except(AttributeError):
                print("AttributeError check render")

    pygame.display.flip()


# this is a script to read gloabl mean temperatures from a csv file and then construct a nice looking time series
import csv
import pandas as pd
import numpy as np

file_name="/home/rusul/Desktop/Upload/global-temp-annual_csv.csv"

#with open(file_name, 'r') as csvfile:
    # creating a csv reader object
 #   csvreader = csv.reader(csvfile, delimiter=',')

  #  for line in csvreader:
   #     print(line[0])

# pandas is better!

df = pd.read_csv(file_name)

df.head() # what data do we want?

temperature=df['Land and Ocean']

temp=np.array(temperature)
temp.shape

# Now let us plot the data

fig= plt.figure(figsize=(6,3))

plt.plot(np.arange(0,len(temp))+1880,temp)
plt.xlabel('year', size=14)
plt.ylabel('anomalies from the mean', size=14)
plt.title('The rise in global mean temperatures since 1880, a clear trend!',size=14)
plt.show()


"""
Serializers from financialaid
"""
import datetime

from django.shortcuts import get_object_or_404
from rest_framework import serializers
from rest_framework.exceptions import ValidationError
from rest_framework.fields import IntegerField, FloatField, CharField

from courses.models import Program
from dashboard.models import ProgramEnrollment
from financialaid.api import determine_tier_program, determine_auto_approval
from financialaid.models import FinancialAid, FinancialAidStatus

class FinancialAidSerializer(serializers.Serializer):
    """
    Serializer for Financial Aid objects
    """
    original_income = FloatField(min_value=0)
    original_currency = CharField()
    program_id = IntegerField()

    def validate(self, data):
        """
        Validators for this serializer
        """
        data["program"] = get_object_or_404(Program, pk=data["program_id"])
        if not data["program"].financial_aid_availability:
            raise ValidationError("Financial aid not available for this program.")
        if not ProgramEnrollment.objects.filter(program=data["program"], user=self.context["request"].user).exists():
            raise ValidationError("User not in program.")
        return data

    def save(self):
        """
        Override save method
        """
        if self.validated_data["original_currency"] != "USD":
            raise ValidationError("Only USD supported currently")
        user = self.context["request"].user
        tier_program = determine_tier_program(self.validated_data["program"], self.validated_data["original_income"])

        financial_aid = FinancialAid.objects.create(
            original_income=self.validated_data["original_income"],
            original_currency=self.validated_data["original_currency"],
            tier_program=tier_program,
            user=user,
            income_usd=self.validated_data["original_income"],
            country_of_income=user.profile.country,
            date_exchange_rate=datetime.datetime.now()
        )

        if determine_auto_approval(financial_aid) is True:
            financial_aid.status = FinancialAidStatus.AUTO_APPROVED
        else:
            financial_aid.status = FinancialAidStatus.PENDING_DOCS
        financial_aid.save()

        # Add auditing here

        return financial_aid