import argparse
import logging
import math
import os
import random
import time

import numpy as np
import torch
from torch.cuda import amp
from torch import nn
from torch.nn import functional as F
from torch import optim
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from data import DATASET_GETTERS
from models import WideResNet, ModelEMA
from utils import (AverageMeter, accuracy, create_loss_fn,
                   save_checkpoint, reduce_tensor, model_load_state_dict)

logger = logging.getLogger(__name__)

parser = argparse.ArgumentParser()
parser.add_argument('--name', type=str, required=True, help='experiment name')
parser.add_argument('--data-path', default='./data', type=str, help='data path')
parser.add_argument('--save-path', default='./checkpoint', type=str, help='save path')
parser.add_argument('--dataset', default='cifar10', type=str,
                    choices=['cifar10', 'cifar100'], help='dataset name')
parser.add_argument('--num-labeled', type=int, default=4000, help='number of labeled data')
parser.add_argument("--expand-labels", action="store_true", help="expand labels to fit eval steps")
parser.add_argument('--total-steps', default=300000, type=int, help='number of total steps to run')
parser.add_argument('--eval-step', default=1000, type=int, help='number of eval steps to run')
parser.add_argument('--start-step', default=0, type=int,
                    help='manual epoch number (useful on restarts)')
parser.add_argument('--workers', default=4, type=int, help='number of workers')
parser.add_argument('--num-classes', default=10, type=int, help='number of classes')
parser.add_argument('--resize', default=32, type=int, help='resize image')
parser.add_argument('--batch-size', default=64, type=int, help='train batch size')
parser.add_argument('--teacher-dropout', default=0, type=float, help='dropout on last dense layer')
parser.add_argument('--student-dropout', default=0, type=float, help='dropout on last dense layer')
parser.add_argument('--teacher_lr', default=0.01, type=float, help='train learning late')
parser.add_argument('--student_lr', default=0.01, type=float, help='train learning late')
parser.add_argument('--momentum', default=0.9, type=float, help='SGD Momentum')
parser.add_argument('--nesterov', action='store_true', help='use nesterov')
parser.add_argument('--weight-decay', default=0, type=float, help='train weight decay')
parser.add_argument('--ema', default=0, type=float, help='EMA decay rate')
parser.add_argument('--warmup-steps', default=0, type=int, help='warmup steps')
parser.add_argument('--student-wait-steps', default=0, type=int, help='warmup steps')
parser.add_argument('--grad-clip', default=0., type=float, help='gradient norm clipping')
parser.add_argument('--resume', default='', type=str, help='path to checkpoint')
parser.add_argument('--evaluate', action='store_true', help='only evaluate model on validation set')
parser.add_argument('--finetune', action='store_true',
                    help='only finetune model on labeled dataset')
parser.add_argument('--finetune-epochs', default=125, type=int, help='finetune epochs')
parser.add_argument('--finetune-batch-size', default=512, type=int, help='finetune batch size')
parser.add_argument('--finetune-lr', default=1e-5, type=float, help='finetune learning late')
parser.add_argument('--finetune-weight-decay', default=0, type=float, help='finetune weight decay')
parser.add_argument('--finetune-momentum', default=0, type=float, help='finetune SGD Momentum')
parser.add_argument('--seed', default=None, type=int, help='seed for initializing training')
parser.add_argument('--label-smoothing', default=0, type=float, help='label smoothing alpha')
parser.add_argument('--mu', default=7, type=int, help='coefficient of unlabeled batch size')
parser.add_argument('--threshold', default=0.95, type=float, help='pseudo label threshold')
parser.add_argument('--temperature', default=1, type=float, help='pseudo label temperature')
parser.add_argument('--lambda-u', default=1, type=float, help='coefficient of unlabeled loss')
parser.add_argument('--uda-steps', default=1, type=float, help='warmup steps of lambda-u')
parser.add_argument("--randaug", nargs="+", type=int, help="use it like this. --randaug 2 10")
parser.add_argument("--amp", action="store_true", help="use 16-bit (mixed) precision")
parser.add_argument('--world-size', default=-1, type=int,
                    help='number of nodes for distributed training')
parser.add_argument("--local_rank", type=int, default=-1,
                    help="For distributed training: local_rank")

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

def get_cosine_schedule_with_warmup(optimizer,
                                    num_warmup_steps,
                                    num_training_steps,
                                    num_wait_steps=0,
                                    num_cycles=0.5,
                                    last_epoch=-1):
    def lr_lambda(current_step):
        if current_step < num_wait_steps:
            return 0.0

        if current_step < num_warmup_steps + num_wait_steps:
            return float(current_step) / float(max(1, num_warmup_steps + num_wait_steps))

        progress = float(current_step - num_warmup_steps - num_wait_steps) / \
            float(max(1, num_training_steps - num_warmup_steps - num_wait_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))

    return LambdaLR(optimizer, lr_lambda, last_epoch)

def get_lr(optimizer):
    return optimizer.param_groups[0]['lr']

def train_loop(args, labeled_loader, unlabeled_loader, test_loader,
               teacher_model, student_model, avg_student_model, criterion,
               t_optimizer, s_optimizer, t_scheduler, s_scheduler, t_scaler, s_scaler):
    logger.info("***** Running Training *****")
    logger.info(f"   Task = {args.dataset}@{args.num_labeled}")
    logger.info(f"   Total steps = {args.total_steps}")

    if args.world_size > 1:
        labeled_epoch = 0
        unlabeled_epoch = 0
        labeled_loader.sampler.set_epoch(labeled_epoch)
        unlabeled_loader.sampler.set_epoch(unlabeled_epoch)

    labeled_iter = iter(labeled_loader)
    unlabeled_iter = iter(unlabeled_loader)

#     moving_dot_product = torch.empty(1).to(args.device)
#     limit = 3.0**(0.5)  # 3 = 6 / (f_in + f_out)
#     nn.init.uniform_(moving_dot_product, -limit, limit)

    for step in range(args.start_step, args.total_steps):
        if step % args.eval_step == 0:
            pbar = tqdm(range(args.eval_step), disable=args.local_rank not in [-1, 0])
            batch_time = AverageMeter()
            data_time = AverageMeter()
            s_losses = AverageMeter()
            t_losses = AverageMeter()
            t_losses_l = AverageMeter()
            t_losses_u = AverageMeter()
            t_losses_mpl = AverageMeter()
            mean_mask = AverageMeter()

        teacher_model.train()
        student_model.train()
        end = time.time()

        try:
            images_l, targets = labeled_iter.next()
        except:
            if args.world_size > 1:
                labeled_epoch += 1
                labeled_loader.sampler.set_epoch(labeled_epoch)
            labeled_iter = iter(labeled_loader)
            images_l, targets = labeled_iter.next()

        try:
            (images_uw, images_us), _ = unlabeled_iter.next()
        except:
            if args.world_size > 1:
                unlabeled_epoch += 1
                unlabeled_loader.sampler.set_epoch(unlabeled_epoch)
            unlabeled_iter = iter(unlabeled_loader)
            (images_uw, images_us), _ = unlabeled_iter.next()

        data_time.update(time.time() - end)

        images_l = images_l.to(args.device)
        images_uw = images_uw.to(args.device)
        images_us = images_us.to(args.device)
        targets = targets.to(args.device)
        with amp.autocast(enabled=args.amp):
            batch_size = images_l.shape[0]
            t_images = torch.cat((images_l, images_uw, images_us))
            t_logits = teacher_model(t_images)
            t_logits_l = t_logits[:batch_size]
            t_logits_uw, t_logits_us = t_logits[batch_size:].chunk(2)
            del t_logits

            t_loss_l = criterion(t_logits_l, targets)

            soft_pseudo_label = torch.softmax(t_logits_uw.detach()/args.temperature, dim=-1)
            max_probs, hard_pseudo_label = torch.max(soft_pseudo_label, dim=-1)
            mask = max_probs.ge(args.threshold).float()
            t_loss_u = torch.mean(
                -(soft_pseudo_label * torch.log_softmax(t_logits_us, dim=-1)).sum(dim=-1) * mask
            )
            weight_u = args.lambda_u * min(1., (step+1) / args.uda_steps)
            t_loss_uda = t_loss_l + weight_u * t_loss_u

            s_images = torch.cat((images_l, images_us))
            s_logits = student_model(s_images)
            s_logits_l = s_logits[:batch_size]
            s_logits_us = s_logits[batch_size:]
            del s_logits

            s_loss_l_old = F.cross_entropy(s_logits_l.detach(), targets)
            s_loss = criterion(s_logits_us, hard_pseudo_label)

        s_scaler.scale(s_loss).backward()
        if args.grad_clip > 0:
            s_scaler.unscale_(s_optimizer)
            nn.utils.clip_grad_norm_(student_model.parameters(), args.grad_clip)
        s_scaler.step(s_optimizer)
        s_scaler.update()
        s_scheduler.step()
        if args.ema > 0:
            avg_student_model.update_parameters(student_model)

        with amp.autocast(enabled=args.amp):
            with torch.no_grad():
                s_logits_l = student_model(images_l)
            s_loss_l_new = F.cross_entropy(s_logits_l.detach(), targets)
#             dot_product = s_loss_l_new - s_loss_l_old
            # test
            dot_product = s_loss_l_old - s_loss_l_new
#             moving_dot_product = moving_dot_product * 0.99 + dot_product * 0.01
#             dot_product = dot_product - moving_dot_product
            _, hard_pseudo_label = torch.max(t_logits_us.detach(), dim=-1)
            t_loss_mpl = dot_product * F.cross_entropy(t_logits_us, hard_pseudo_label)
            t_loss = t_loss_uda + t_loss_mpl

        t_scaler.scale(t_loss).backward()
        if args.grad_clip > 0:
            t_scaler.unscale_(t_optimizer)
            nn.utils.clip_grad_norm_(teacher_model.parameters(), args.grad_clip)
        t_scaler.step(t_optimizer)
        t_scaler.update()
        t_scheduler.step()

        teacher_model.zero_grad()
        student_model.zero_grad()

        if args.world_size > 1:
            s_loss = reduce_tensor(s_loss.detach(), args.world_size)
            t_loss = reduce_tensor(t_loss.detach(), args.world_size)
            t_loss_l = reduce_tensor(t_loss_l.detach(), args.world_size)
            t_loss_u = reduce_tensor(t_loss_u.detach(), args.world_size)
            t_loss_mpl = reduce_tensor(t_loss_mpl.detach(), args.world_size)
            mask = reduce_tensor(mask, args.world_size)

        s_losses.update(s_loss.item())
        t_losses.update(t_loss.item())
        t_losses_l.update(t_loss_l.item())
        t_losses_u.update(t_loss_u.item())
        t_losses_mpl.update(t_loss_mpl.item())
        mean_mask.update(mask.mean().item())

        batch_time.update(time.time() - end)
        pbar.set_description(
            f"Train Iter: {step+1:3}/{args.total_steps:3}. "
            f"LR: {get_lr(s_optimizer):.4f}. Data: {data_time.avg:.2f}s. "
            f"Batch: {batch_time.avg:.2f}s. S_Loss: {s_losses.avg:.4f}. "
            f"T_Loss: {t_losses.avg:.4f}. Mask: {mean_mask.avg:.4f}. ")
        pbar.update()
        if args.local_rank in [-1, 0]:
            args.writer.add_scalar("lr", get_lr(s_optimizer), step)

        args.num_eval = step//args.eval_step
        if (step+1) % args.eval_step == 0:
            pbar.close()
            if args.local_rank in [-1, 0]:
                args.writer.add_scalar("train/1.s_loss", s_losses.avg, args.num_eval)
                args.writer.add_scalar("train/2.t_loss", t_losses.avg, args.num_eval)
                args.writer.add_scalar("train/3.t_labeled", t_losses_l.avg, args.num_eval)
                args.writer.add_scalar("train/4.t_unlabeled", t_losses_u.avg, args.num_eval)
                args.writer.add_scalar("train/5.t_mpl", t_losses_mpl.avg, args.num_eval)
                args.writer.add_scalar("train/6.mask", mean_mask.avg, args.num_eval)

                test_model = avg_student_model if avg_student_model is not None else student_model
                test_loss, top1, top5 = evaluate(args, test_loader, test_model, criterion)

                args.writer.add_scalar("test/loss", test_loss, args.num_eval)
                args.writer.add_scalar("test/acc@1", top1, args.num_eval)
                args.writer.add_scalar("test/acc@5", top5, args.num_eval)

                is_best = top1 > args.best_top1
                if is_best:
                    args.best_top1 = top1
                    args.best_top5 = top5

                logger.info(f"top-1 acc: {top1:.2f}")
                logger.info(f"Best top-1 acc: {args.best_top1:.2f}")

                save_checkpoint(args, {
                    'step': step + 1,
                    'teacher_state_dict': teacher_model.state_dict(),
                    'student_state_dict': student_model.state_dict(),
                    'avg_state_dict': avg_student_model.state_dict() if avg_student_model is not None else None,
                    'best_top1': args.best_top1,
                    'best_top5': args.best_top5,
                    'teacher_optimizer': t_optimizer.state_dict(),
                    'student_optimizer': s_optimizer.state_dict(),
                    'teacher_scheduler': t_scheduler.state_dict(),
                    'student_scheduler': s_scheduler.state_dict(),
                    'teacher_scaler': t_scaler.state_dict(),
                    'student_scaler': s_scaler.state_dict(),
                }, is_best)
    # finetune
    del t_scaler, t_scheduler, t_optimizer, teacher_model, unlabeled_loader
    del s_scaler, s_scheduler, s_optimizer
    ckpt_name = f'{args.save_path}/{args.name}_best.pth.tar'
    loc = f'cuda:{args.gpu}'
    checkpoint = torch.load(ckpt_name, map_location=loc)
    logger.info(f"=> loading checkpoint '{ckpt_name}'")
    if checkpoint['avg_state_dict'] is not None:
        model_load_state_dict(student_model, checkpoint['avg_state_dict'])
    else:
        model_load_state_dict(student_model, checkpoint['student_state_dict'])
    finetune(args, labeled_loader, test_loader, student_model, criterion)
    return

def evaluate(args, test_loader, model, criterion):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()
    model.eval()
    test_iter = tqdm(test_loader, disable=args.local_rank not in [-1, 0])
    with torch.no_grad():
        end = time.time()
        for step, (images, targets) in enumerate(test_iter):
            data_time.update(time.time() - end)
            batch_size = targets.shape[0]
            images = images.to(args.device)
            targets = targets.to(args.device)
            with amp.autocast(enabled=args.amp):
                outputs = model(images)
                loss = criterion(outputs, targets)

            acc1, acc5 = accuracy(outputs, targets, (1, 5))
            losses.update(loss.item(), batch_size)
            top1.update(acc1[0], batch_size)
            top5.update(acc5[0], batch_size)
            batch_time.update(time.time() - end)
            end = time.time()
            test_iter.set_description(
                f"Test Iter: {step+1:3}/{len(test_loader):3}. Data: {data_time.avg:.2f}s. "
                f"Batch: {batch_time.avg:.2f}s. Loss: {losses.avg:.4f}. "
                f"top1: {top1.avg:.2f}. top5: {top5.avg:.2f}. ")

        test_iter.close()
        return losses.avg, top1.avg, top5.avg

def finetune(args, train_loader, test_loader, model, criterion):
    train_sampler = RandomSampler if args.local_rank == -1 else DistributedSampler
    labeled_loader = DataLoader(
        train_loader.dataset,
        sampler=train_sampler(train_loader.dataset),
        batch_size=args.finetune_batch_size,
        num_workers=args.workers,
        pin_memory=True)
    optimizer = optim.SGD(model.parameters(),
                          lr=args.finetune_lr,
                          momentum=args.finetune_momentum,
                          weight_decay=args.finetune_weight_decay)
    scaler = amp.GradScaler(enabled=args.amp)

    logger.info("***** Running Finetuning *****")
    logger.info(f"   Finetuning steps = {len(labeled_loader)*args.finetune_epochs}")

    for epoch in range(args.finetune_epochs):
        if args.world_size > 1:
            labeled_loader.sampler.set_epoch(epoch+624)

        batch_time = AverageMeter()
        data_time = AverageMeter()
        losses = AverageMeter()
        model.train()
        end = time.time()
        labeled_iter = tqdm(labeled_loader, disable=args.local_rank not in [-1, 0])
        for step, (images, targets) in enumerate(labeled_iter):
            data_time.update(time.time() - end)
            batch_size = targets.shape[0]
            images = images.to(args.device)
            targets = targets.to(args.device)
            with amp.autocast(enabled=args.amp):
                model.zero_grad()
                outputs = model(images)
                loss = criterion(outputs, targets)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            if args.world_size > 1:
                loss = reduce_tensor(loss.detach(), args.world_size)
            losses.update(loss.item(), batch_size)
            batch_time.update(time.time() - end)
            labeled_iter.set_description(
                f"Finetune Epoch: {epoch+1:2}/{args.finetune_epochs:2}. Data: {data_time.avg:.2f}s. "
                f"Batch: {batch_time.avg:.2f}s. Loss: {losses.avg:.4f}. ")
        labeled_iter.close()
        if args.local_rank in [-1, 0]:
            args.writer.add_scalar("finetune/train_loss", losses.avg, epoch)
            test_loss, top1, top5 = evaluate(args, test_loader, model, criterion)
            args.writer.add_scalar("finetune/test_loss", test_loss, epoch)
            args.writer.add_scalar("finetune/acc@1", top1, epoch)
            args.writer.add_scalar("finetune/acc@5", top5, epoch)
            is_best = top1 > args.best_top1
            if is_best:
                args.best_top1 = top1
                args.best_top5 = top5

            logger.info(f"top-1 acc: {top1:.2f}")
            logger.info(f"Best top-1 acc: {args.best_top1:.2f}")

            save_checkpoint(args, {
                'step': step + 1,
                'best_top1': args.best_top1,
                'best_top5': args.best_top5,
                'student_state_dict': model.state_dict(),
                'avg_state_dict': None,
                'student_optimizer': optimizer.state_dict(),
            }, is_best, finetune=True)
    return

def main():
    args = parser.parse_args()
    args.best_top1 = 0.
    args.best_top5 = 0.

    if args.local_rank != -1:
        args.gpu = args.local_rank
        torch.distributed.init_process_group(backend='nccl')
        args.world_size = torch.distributed.get_world_size()
    else:
        args.gpu = 0
        args.world_size = 1

    args.device = torch.device('cuda', args.gpu)

    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARNING)

    logger.warning(
        f"Process rank: {args.local_rank}, "
        f"device: {args.device}, "
        f"distributed training: {bool(args.local_rank != -1)}, "
        f"16-bits training: {args.amp}")

    logger.info(dict(args._get_kwargs()))

    if args.local_rank in [-1, 0]:
        args.writer = SummaryWriter(f"results/{args.name}")

    if args.seed is not None:
        set_seed(args)

    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()

    labeled_dataset, unlabeled_dataset, test_dataset = DATASET_GETTERS[args.dataset](args)

    if args.local_rank == 0:
        torch.distributed.barrier()

    train_sampler = RandomSampler if args.local_rank == -1 else DistributedSampler
    labeled_loader = DataLoader(
        labeled_dataset,
        sampler=train_sampler(labeled_dataset),
        batch_size=args.batch_size,
        num_workers=args.workers,
        drop_last=True)

    unlabeled_loader = DataLoader(
        unlabeled_dataset,
        sampler=train_sampler(unlabeled_dataset),
        batch_size=args.batch_size*args.mu,
        num_workers=args.workers,
        drop_last=True)

    test_loader = DataLoader(test_dataset,
                             sampler=SequentialSampler(test_dataset),
                             batch_size=args.batch_size,
                             num_workers=args.workers)

    if args.dataset == "cifar10":
        depth, widen_factor = 28, 2
    elif args.dataset == 'cifar100':
        depth, widen_factor = 28, 8

    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()

    teacher_model = WideResNet(num_classes=args.num_classes,
                               depth=depth,
                               widen_factor=widen_factor,
                               dropout=0,
                               dense_dropout=args.teacher_dropout)
    student_model = WideResNet(num_classes=args.num_classes,
                               depth=depth,
                               widen_factor=widen_factor,
                               dropout=0,
                               dense_dropout=args.student_dropout)

    if args.local_rank == 0:
        torch.distributed.barrier()

    logger.info(f"Model: WideResNet {depth}x{widen_factor}")
    logger.info(f"Params: {sum(p.numel() for p in teacher_model.parameters())/1e6:.2f}M")

    teacher_model.to(args.device)
    student_model.to(args.device)
    avg_student_model = None
    if args.ema > 0:
        avg_student_model = ModelEMA(student_model, args.ema)

    criterion = create_loss_fn(args)

    no_decay = ['bn']
    teacher_parameters = [
        {'params': [p for n, p in teacher_model.named_parameters() if not any(
            nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in teacher_model.named_parameters() if any(
            nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]
    student_parameters = [
        {'params': [p for n, p in student_model.named_parameters() if not any(
            nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in student_model.named_parameters() if any(
            nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]

    t_optimizer = optim.SGD(teacher_parameters,
                            lr=args.teacher_lr,
                            momentum=args.momentum,
                            # weight_decay=args.weight_decay,
                            nesterov=args.nesterov)
    s_optimizer = optim.SGD(student_parameters,
                            lr=args.student_lr,
                            momentum=args.momentum,
                            # weight_decay=args.weight_decay,
                            nesterov=args.nesterov)

    t_scheduler = get_cosine_schedule_with_warmup(t_optimizer,
                                                  args.warmup_steps,
                                                  args.total_steps)
    s_scheduler = get_cosine_schedule_with_warmup(s_optimizer,
                                                  args.warmup_steps,
                                                  args.total_steps,
                                                  args.student_wait_steps)

    t_scaler = amp.GradScaler(enabled=args.amp)
    s_scaler = amp.GradScaler(enabled=args.amp)

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            logger.info(f"=> loading checkpoint '{args.resume}'")
            loc = f'cuda:{args.gpu}'
            checkpoint = torch.load(args.resume, map_location=loc)
            args.best_top1 = checkpoint['best_top1'].to(torch.device('cpu'))
            args.best_top5 = checkpoint['best_top5'].to(torch.device('cpu'))
            if not (args.evaluate or args.finetune):
                args.start_step = checkpoint['step']
                t_optimizer.load_state_dict(checkpoint['teacher_optimizer'])
                s_optimizer.load_state_dict(checkpoint['student_optimizer'])
                t_scheduler.load_state_dict(checkpoint['teacher_scheduler'])
                s_scheduler.load_state_dict(checkpoint['student_scheduler'])
                t_scaler.load_state_dict(checkpoint['teacher_scaler'])
                s_scaler.load_state_dict(checkpoint['student_scaler'])
                model_load_state_dict(teacher_model, checkpoint['teacher_state_dict'])
                if avg_student_model is not None:
                    model_load_state_dict(avg_student_model, checkpoint['avg_state_dict'])

            else:
                if checkpoint['avg_state_dict'] is not None:
                    model_load_state_dict(student_model, checkpoint['avg_state_dict'])
                else:
                    model_load_state_dict(student_model, checkpoint['student_state_dict'])

            logger.info(f"=> loaded checkpoint '{args.resume}' (step {checkpoint['step']})")
        else:
            logger.info(f"=> no checkpoint found at '{args.resume}'")

    if args.local_rank != -1:
        teacher_model = nn.parallel.DistributedDataParallel(
            teacher_model, device_ids=[args.local_rank],
            output_device=args.local_rank, find_unused_parameters=True)
        student_model = nn.parallel.DistributedDataParallel(
            student_model, device_ids=[args.local_rank],
            output_device=args.local_rank, find_unused_parameters=True)

    if args.finetune:
        del t_scaler, t_scheduler, t_optimizer, teacher_model, unlabeled_loader
        del s_scaler, s_scheduler, s_optimizer
        finetune(args, labeled_loader, test_loader, student_model, criterion)
        return

    if args.evaluate:
        del t_scaler, t_scheduler, t_optimizer, teacher_model, unlabeled_loader, labeled_loader
        del s_scaler, s_scheduler, s_optimizer
        evaluate(args, test_loader, student_model, criterion)
        return

    teacher_model.zero_grad()
    student_model.zero_grad()
    train_loop(args, labeled_loader, unlabeled_loader, test_loader,
               teacher_model, student_model, avg_student_model, criterion,
               t_optimizer, s_optimizer, t_scheduler, s_scheduler, t_scaler, s_scaler)
    return

if __name__ == '__main__':
    main()


import copy
import math
import pickle as pkl
import numpy as np
import torch
import os
import sys
import torch.nn as nn
import torch.nn.functional as F
from gym import spaces

from ...util.model import BenchmarkModel
from torchbenchmark.tasks import REINFORCEMENT_LEARNING

from .utils import FrameStack, set_seed_everywhere, eval_mode
from .drq import DRQAgent
from .config import DRQConfig
from .replay_buffer import ReplayBuffer

class MockEnv:
    def __init__(self, obs):
        self._norm_action_space = spaces.Box(
            low=-1.0,
            high=1.0,
            shape=[1],
            dtype=np.float32)
        self._observation_space = spaces.Box(
            low=0,
            high=255,
            shape=[9, 84, 84],
            dtype=np.uint8
        )
        self.obs = obs
        self._max_episode_steps = 250
        self.metadata = {'render.modes': []}
        self.reward_range = (-float('inf'), float('inf'))
    def step(self, action):
        reward = 0.0
        done = False
        info_state = [0.016243, 3.1355, -0.0052817, -0.01073]
        info = dict()
        info["internal_state"] = info_state
        info["discount"] = 1.0
        return (self.obs, reward, done, info)
    def seed(self, seed=None):
        self._norm_action_space.seed(seed)
        self._observation_space.seed(seed)
    def reset(self):
        return self.obs
    @property
    def observation_space(self):
        return self._observation_space
    @property
    def action_space(self):
        return self._norm_action_space

def make_env(cfg):
    if cfg.env == 'ball_in_cup_catch':
        domain_name = 'ball_in_cup'
        task_name = 'catch'
    elif cfg.env == 'point_mass_easy':
        domain_name = 'point_mass'
        task_name = 'easy'
    else:
        domain_name = cfg.env.split('_')[0]
        task_name = '_'.join(cfg.env.split('_')[1:])
     # per dreamer: https://github.com/danijar/dreamer/blob/02f0210f5991c7710826ca7881f19c64a012290c/wrappers.py#L26
    camera_id = 2 if domain_name == 'quadruped' else 0

    current_dir = os.path.dirname(os.path.realpath(__file__))
    mockobs = pkl.load(open(os.path.join(current_dir, cfg.obs_path), "rb"))
    low = np.amin(mockobs)
    high = np.amax(mockobs)
    mockobs = np.random.randint(low=11, high=228, size=mockobs.shape, dtype=np.uint8)
    env = MockEnv(mockobs)
    env = FrameStack(env, k=cfg.frame_stack)

    env.seed(cfg.seed)
    assert env.action_space.low.min() >= -1
    assert env.action_space.high.max() <= 1

    return env

class Model(BenchmarkModel):
    task = REINFORCEMENT_LEARNING.OTHER_RL
    def __init__(self, device=None, jit=False):
        super(Model, self).__init__()
        self.device = device
        self.jit = jit
        self.cfg = DRQConfig()
        set_seed_everywhere(self.cfg.seed)
        self.env = make_env(self.cfg)
        obs_shape = self.env.observation_space.shape
        action_shape = self.env.action_space.shape
        action_range = [
            float(self.env.action_space.low.min()),
            float(self.env.action_space.high.max())
        ]
        self.agent = DRQAgent(self.cfg, self.device, obs_shape, action_shape, action_range)
        self.replay_buffer = ReplayBuffer(self.env.observation_space.shape,
                                          self.env.action_space.shape,
                                          self.cfg.replay_buffer_capacity,
                                          self.cfg.image_pad, self.device)
        self.step = 0

    def get_module(self):
        obs = self.env.reset()
        obs = torch.FloatTensor(obs).to(self.device)
        obs = obs.unsqueeze(0)
        return self.agent.actor, (obs, )

    def train(self, niter=2):
        if self.jit:
            raise NotImplementedError()
        episode, episode_reward, episode_step, done = 0, 0, 1, True
        for step in range(niter):
            obs = self.env.reset()
            done = False
            episode_reward = 0
            episode_step = 0
            episode += 1
            if step < self.cfg.num_seed_steps:
                action = self.env.action_space.sample()
            else:
                with eval_mode(self.agent):
                    action = self.agent.act(obs, sample=True)
            # run training update
            if self.step >= self.cfg.num_seed_steps:
                for _ in range(self.cfg.num_train_iters):
                    self.agent.update(self.replay_buffer, None,
                                      self.step)
            next_obs, reward, done, info = self.env.step(action)
            # allow infinite bootstrap
            done = float(done)
            done_no_max = 0 if episode_step + 1 == self.env._max_episode_steps else done
            episode_reward += reward
            self.replay_buffer.add(obs, action, reward, next_obs, done,
                                   done_no_max)
            obs = next_obs
            episode_step += 1
            self.step += 1

    def eval(self, niter=1):
        if self.jit:
            raise NotImplementedError()
        average_episode_reward = 0
        for episode in range(niter):
            obs = self.env.reset()
            episode_reward = 0
            episode_step = 0
            with eval_mode(self.agent):
                action = self.agent.act(obs, sample=False)
            obs, reward, done, info = self.env.step(action)
            episode_reward += reward
            episode_step += 1
            average_episode_reward += episode_reward
        average_episode_reward /= float(niter)


"""
Metabolite annotation
"""
import numpy as np
from medipy.io import load
import os
import fonc_util as f
import gen_in
import lawij
import ser_xml
import xml_gen
import locale

def ann(root,wa7da):
    old_locale = locale.getlocale()
    locale.setlocale(locale.LC_ALL, "C")
    tabp,tabc=gen_in.ini1(root)

    boc=-1
    lolo= load(os.path.join(root,"2rr"))
    H=lolo.data
    D=H[0,:,:]
    H=D
    #print corps.interquartile_range(H)
    #print np.size(H,0),np.size(H,1)
    list_exp=ser_xml.exceptions(root)
    raxep=[]
    while boc<len(tabp)-1:
        boc+=1
        a=str(tabp[boc][2])
        #print a,boc
        newn=''
        for j in range(len(a)):
            if (a[j]=='_')==1:
                break
            newn+=a[j]

        r1=[]
        r1.append((boc))
        for jj in range(boc+1,len(tabp)):
            nomn=str(tabp[jj][2])
            #print nomn[0:j]
            try:
                if nomn[0:j+1]==newn+'_':
                    r1.append((jj))
                    #print 'ok'
                else:
                    break
            except:
                    break
        boc=boc+len(r1)-1
        #print tabp[r1[0]:r1[len(r1)-1]+1]
        #raw_input()
        #print len(r1)
        nt=tabp[r1[0]:r1[len(r1)-1]+1]
        #print nt
        #print 'Start'
        test=[]
        testc3=[]
        con=0
        ampref=[]
        amp=[]
        newrt=[]
        #print newn
        ham=0
        ed5il=0
        for jj in range(len(nt)):

            #print newn
            #print jj
            #print nt[jj][0],nt[jj][1]
            #print nt[jj][0],nt[jj][1]
            r,indi=lawij.alig(nt[jj][0],nt[jj][1],H,nt[jj][3],wa7da)

            #print r,indi,nt[jj][2],nt[jj][0]-r[0],nt[jj][1]-r[1],nt[jj][7]

            if nt[jj][7]=='n':
                ed5il=1

            #raw_input()
            if r[0]==100 and nt[jj][7]=='y':
                #print r,indi,nt[jj][2],nt[jj][0]-r[0],nt[jj][1]-r[1],nt[jj][7]
                ham=1
                break

            if indi==0 :
                con=con+1

                if np.abs(r[0])==4 or np.abs(r[1])==6:
                    testc3.append((1))

                else:
                    testc3.append((0))
                test.append((0))
                #fig = plt.figure()
                #ax = fig.add_subplot(111)
                zayneb=H[(nt[jj][0]-r[0])-3:(nt[jj][0]-r[0])+4,(nt[jj][1]-r[1])-3:(nt[jj][1]-r[1])+4]*1.
                nr=f.subpix2(zayneb,(nt[jj][0]-r[0]),(nt[jj][1]-r[1]))
                #print (nt[jj][0]-r[0]),(nt[jj][1]-r[1]),nr
                zayneb=H[(nt[jj][0]-r[0])-9:(nt[jj][0]-r[0])+10,(nt[jj][1]-r[1])-9:(nt[jj][1]-r[1])+10]*1.
                #nr=f.subpix(zayneb,(nt[jj][0]-r[0]),(nt[jj][1]-r[1]))
                #print (nt[jj][0]-r[0]),(nt[jj][1]-r[1]),nr
                chl=f.dephcl(zayneb)
                chg=f.dephcg(zayneb)
                #ch,congl=f.dephc(zayneb)
                ch,congl=f.dephcaprio(zayneb,float(nt[jj][4]),float(nt[jj][5]),nt[jj][6])
                #print congl
                #cax = ax.imshow(zayneb, interpolation='nearest')
                #plt.show()
                #print ch
                #plt.show()
                #print nt[jj][3]
                ampref.append(float(nt[jj][3])*1.)
                amp.append( H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])]*1.)
                #print str(H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])])
                newrt.append((nr[0],nr[1],nt[jj][2],str(H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])]),float(nt[jj][3]),chg[1],chl[0],chl[1],ch[0],ch[1],congl,r[0],r[1]))
                #print newrt
            else:
                if r[0]<100 :
                    if np.abs(r[0])==3 or np.abs(r[1])==5:
                        testc3.append((1))

                    else:
                        testc3.append((0))
                    con=con+1
                    test.append((1))
                    #fig = plt.figure()
                    #ax = fig.add_subplot(111)
                    zayneb=H[(nt[jj][0]-r[0])-3:(nt[jj][0]-r[0])+4,(nt[jj][1]-r[1])-3:(nt[jj][1]-r[1])+4]*1.
                    nr=f.subpix2(zayneb,(nt[jj][0]-r[0]),(nt[jj][1]-r[1]))
                    #print (nt[jj][0]-r[0]),(nt[jj][1]-r[1]),nr
                    zayneb=H[(nt[jj][0]-r[0])-9:(nt[jj][0]-r[0])+10,(nt[jj][1]-r[1])-9:(nt[jj][1]-r[1])+10]*1.
                    #nr=f.subpix(zayneb,(nt[jj][0]-r[0]),(nt[jj][1]-r[1]))
                    #print (nt[jj][0]-r[0]),(nt[jj][1]-r[1]),nr
                    chl=f.dephcl(zayneb)
                    chg=f.dephcg(zayneb)
                    #ch,congl=f.dephc(zayneb)
                    ch,congl=f.dephcaprio(zayneb,float(nt[jj][4]),float(nt[jj][5]),nt[jj][6])
                    ampref.append(float(nt[jj][3])*1.)
                    amp.append(H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])]*1.)
                    newrt.append((nr[0],nr[1],nt[jj][2],H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])]*1.,float(nt[jj][3]),chg[1],chl[0],chl[1],ch[0],ch[1],congl,r[0],r[1]))
                else:
                    test.append((2))
                    testc3.append((2))
                    #fig = plt.figure()
                    #ax = fig.add_subplot(111)
                    zayneb=H[nt[jj][0]-7:(nt[jj][0])+8,(nt[jj][1])-7:(nt[jj][1])+8]*1.
                    ampref.append(float(nt[jj][3])*1.)
                    amp.append(0)
                    #cax = ax.imshow(zayneb, interpolation='nearest')
                    #plt.show()
                    #raw_input()
                    newrt.append((0,0,0,0,0,0,0,0,0,0,0,0,0))

        #raw_input()
        #print newn
        #print ampref
        #print amp,newn,testc3
        #print nt
        #print newrt
        #raw_input()
        #amptest=np.nonzero(amp>0)
        o=np.nonzero(testc3==0)
        vamp=np.array(amp)
        ivamref=vamp[o]
        o=np.nonzero(ivamref>210000)
        #print newn,'test'
        #print 'ham',ham
        #if (float(len(o[0]))*1.000000001)/float(len(ivamref)*1.00000001)>0.4 or f.exp_hand(list_exp,newn)==1 or ham==0:
        if ham==0:
            #print 'd5aal'
            if len(nt)==con:
                if len(nt)==1:
                    raxep.append(newrt[0])
                    #print 'accepted'
                else:
                    artestc3=np.array(testc3)
                    olc3=np.nonzero(artestc3==1)
                    if np.size(olc3)>0:
                        if f.exp_ync(amp,ampref,testc3)==1:
                            #print 'ouffffff'
                            artest=np.array(test)
                            ol=np.nonzero(artest==1)
                            o=np.nonzero(artest==0)
                            if np.size(ol)>0:
                                #print 'accepted with some conditions'
                                #print f.exp_ync(amp,ampref,testc3)
                                #if f.exp_ync(amp,ampref,testc3)==0:
                                    #print 'llllllaaaaaaaaa'
                                if f.exp_yn(amp,ampref,test)==1:
                                    for no in range(len(newrt)):
                                        raxep.append(newrt[no])
                                elif f.exp_hand(list_exp,newn)==1:
                                    artest=np.array(test)
                                    rnt=np.array(newrt)
                                    ol=np.nonzero(artest==1)
                                    o=np.nonzero(artest==0)
                                    vo=rnt[o]
                                    #raw_input()
                                    for no in range(len(vo)):
                                        #print '%%%%%%%%%'
                                        zi=lawij.mod(vo[no])
                                        #print zi
                                        #print artest
                                        #raw_input()
                                        raxep.append((zi))
                            else:
                                #print f.exp_ync(amp,ampref,testc3)
                                for no in range(len(nt)):
                                    raxep.append(newrt[no])
                                #print 'accepted'
                        else:
                            #print 'ouuuuut'
                            #print f.exp_hand(list_exp,newn)
                            if f.exp_hand(list_exp,newn)==1 :
                                artest=np.array(test)
                                artestc3=np.array(testc3)
                                rnt=np.array(newrt)
                                #print nt
                                #raw_input()
                                ol=np.nonzero(artest==1)
                                condlist=[artest==0,artestc3==0]
                                g=condlist[0]*condlist[1]
                                #print g,artest,artestc3
                                o=np.nonzero(artest==0)
                                #print '%%%hhh%%%%%%'
                                #print rnt
                                #print test
                                vo=rnt[g]
                                #print vo,'%%%hhh%%%%%%'
                                #raw_input()
                                for no in range(len(vo)):
                                            #print '%%%%%%%%%'
                                            zi=lawij.mod(vo[no])
                                            #print zi
                                            #print artest
                                            #raw_input()
                                            raxep.append((zi))
                    else:
                                #print f.exp_ync(amp,ampref,testc3)
                                for no in range(len(nt)):
                                    raxep.append(newrt[no])
                                #print 'accepted'

            else:
                if ham==0:
                    if f.exp_hand(list_exp,newn)==1 or ed5il==1:
                        artest=np.array(test)
                        #print newrt
                        rnt=xml_gen.cell2tab8(newrt)
                        #print nt
                        #raw_input()
                        ol=np.nonzero(artest==1)
                        o=np.nonzero(artest==0)
                        #print '%%%hhh%%%%%%'
                        rnt=np.array(rnt)
                        #print test
                        vo=rnt[o]
                        #print vo
                        #print len(vo)
                        #raw_input()
                        for no in range(len(vo)):
                                    #print '%%%%%%%%%'
                                    zi=lawij.mod(vo[no])
                                    #print zi
                                    #print artest
                                    #raw_input()
                                    raxep.append((zi))
                    #print 'may be...'
                #else:
                    #print 'refused without discussion'
            #print test

            #print 'DONE'
            #print '%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%'
            #print '        '
            #raw_input()
    locale.setlocale(locale.LC_ALL, old_locale)
    return raxep,tabc
    #for lll in range(len(raxep)):
        #print  raxep[lll]

def rec(root):
    old_locale = locale.getlocale()
    locale.setlocale(locale.LC_ALL, "C")
    #root='/home/miv/belghith/Bureau/KAROM/Akram/nmr/RP/5/pdata/1'
    tabp,tabc=gen_in.ini2(root)

    boc=-1
    lolo= load(os.path.join(root,"2rr"))
    H=lolo.data
    D=H[0,:,:]
    H=D
    #print corps.interquartile_range(H)
    #print np.size(H,0),np.size(H,1)
    list_exp=ser_xml.exceptions(root)
    raxep=[]
    while boc<len(tabp)-1:
        boc+=1
        a=str(tabp[boc][2])
        #print a,boc
        newn=''
        for j in range(len(a)):
            if (a[j]=='_')==1:
                break
            newn+=a[j]

        r1=[]
        r1.append((boc))
        for jj in range(boc+1,len(tabp)):
            nomn=str(tabp[jj][2])
            #print nomn[0:j]
            try:
                if nomn[0:j+1]==newn+'_':
                    r1.append((jj))
                    #print 'ok'
                else:
                    break
            except:
                    break
        boc=boc+len(r1)-1
        #print tabp[r1[0]:r1[len(r1)-1]+1]
        #raw_input()
        #print len(r1)
        nt=tabp[r1[0]:r1[len(r1)-1]+1]
        #print nt
        #print 'Start'
        test=[]
        testc3=[]
        con=0
        ampref=[]
        amp=[]
        newrt=[]
        #print newn
        for jj in range(len(nt)):
            #print newn
            #print jj
            #print nt[jj][0],nt[jj][1]
            #print nt[jj][0],nt[jj][1]
            r,indi=lawij.aligrec(nt[jj][0],nt[jj][1],H,nt[jj][3])

            #print r,indi,nt[jj][2],nt[jj][0]-r[0],nt[jj][1]-r[1]
            #raw_input()
            if indi==0 :
                con=con+1

                #print 'ok'
                testc3.append((0))
                test.append((0))
                #fig = plt.figure()
                #ax = fig.add_subplot(111)
                zayneb=H[(nt[jj][0]-r[0])-3:(nt[jj][0]-r[0])+4,(nt[jj][1]-r[1])-3:(nt[jj][1]-r[1])+4]*1.
                nr=f.subpix2(zayneb,(nt[jj][0]-r[0]),(nt[jj][1]-r[1]))
                #print (nt[jj][0]-r[0]),(nt[jj][1]-r[1]),nr
                zayneb=H[(nt[jj][0]-r[0])-9:(nt[jj][0]-r[0])+10,(nt[jj][1]-r[1])-9:(nt[jj][1]-r[1])+10]*1.
                #nr=f.subpix(zayneb,(nt[jj][0]-r[0]),(nt[jj][1]-r[1]))
                #print (nt[jj][0]-r[0]),(nt[jj][1]-r[1]),nr
                chl=f.dephcl(zayneb)
                chg=f.dephcg(zayneb)
                #ch,congl=f.dephc(zayneb)
                ch,congl=f.dephcaprio(zayneb,float(nt[jj][4]),float(nt[jj][5]),nt[jj][6])
                #print congl
                #cax = ax.imshow(zayneb, interpolation='nearest')
                #plt.show()
                #print ch
                #plt.show()
                ampref.append(float(nt[jj][3])*1.)
                amp.append( H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])]*1.)
                #print str(H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])])
                newrt.append((nr[0],nr[1],nt[jj][2],str(H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])]),chg[0],chg[1],chl[0],chl[1],ch[0],ch[1],congl,r[0],r[1]))
                #print newrt
            else:
                if r[0]<100 :
                    if np.abs(r[0])==3 or np.abs(r[1])==5:
                        testc3.append((1))

                    else:
                        testc3.append((0))
                    con=con+1
                    test.append((1))
                    #fig = plt.figure()
                    #ax = fig.add_subplot(111)
                    zayneb=H[(nt[jj][0]-r[0])-3:(nt[jj][0]-r[0])+4,(nt[jj][1]-r[1])-3:(nt[jj][1]-r[1])+4]*1.
                    nr=f.subpix2(zayneb,(nt[jj][0]-r[0]),(nt[jj][1]-r[1]))
                    #print (nt[jj][0]-r[0]),(nt[jj][1]-r[1]),nr
                    zayneb=H[(nt[jj][0]-r[0])-9:(nt[jj][0]-r[0])+10,(nt[jj][1]-r[1])-9:(nt[jj][1]-r[1])+10]*1.
                    #nr=f.subpix(zayneb,(nt[jj][0]-r[0]),(nt[jj][1]-r[1]))
                    #print (nt[jj][0]-r[0]),(nt[jj][1]-r[1]),nr
                    chl=f.dephcl(zayneb)
                    chg=f.dephcg(zayneb)
                    #ch,congl=f.dephc(zayneb)
                    ch,congl=f.dephcaprio(zayneb,float(nt[jj][4]),float(nt[jj][5]),nt[jj][6])
                    ampref.append(float(nt[jj][3])*1.)
                    amp.append(H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])]*1.)
                    newrt.append((nr[0],nr[1],nt[jj][2],H[(nt[jj][0]-r[0]),(nt[jj][1]-r[1])]*1.,chg[0],chg[1],chl[0],chl[1],ch[0],ch[1],congl,r[0],r[1]))
                else:
                    test.append((2))
                    testc3.append((2))
                    #fig = plt.figure()
                    #ax = fig.add_subplot(111)
                    zayneb=H[nt[jj][0]-7:(nt[jj][0])+8,(nt[jj][1])-7:(nt[jj][1])+8]*1.
                    ampref.append(float(nt[jj][3])*1.)
                    amp.append(0)
                    #cax = ax.imshow(zayneb, interpolation='nearest')
                    #plt.show()
                    #raw_input()
                    newrt.append((0,0,0,0,0,0,0,0,0,0,0,0,0))

        #raw_input()
        #print newn
        #print ampref
        #print amp,newn,testc3
        #print nt
        #print newrt
        #raw_input()
        #amptest=np.nonzero(amp>0)
        o=np.nonzero(testc3==0)
        vamp=np.array(amp)
        ivamref=vamp[o]
        o=np.nonzero(ivamref>100)
        #print newn,'test'
        if (float(len(o[0]))*1.000000001)/float(len(ivamref)*1.00000001)>0.4 or f.exp_hand(list_exp,newn)==1:
            #print newn,'d5aal'
            if len(nt)==con:
                if len(nt)==1:
                    raxep.append(newrt[0])
                    #print 'accepted'
                else:
                    artestc3=np.array(testc3)
                    olc3=np.nonzero(artestc3==1)
                    if np.size(olc3)>0:
                        if f.exp_ync(amp,ampref,testc3)==1:
                            #print 'ouffffff'
                            artest=np.array(test)
                            ol=np.nonzero(artest==1)
                            o=np.nonzero(artest==0)
                            if np.size(ol)>0:
                                #print 'accepted with some conditions'
                                #print f.exp_ync(amp,ampref,testc3)
                                #if f.exp_ync(amp,ampref,testc3)==0:
                                    #print 'llllllaaaaaaaaa'
                                if f.exp_yn(amp,ampref,test)==1:
                                    for no in range(len(newrt)):
                                        raxep.append(newrt[no])
                                elif f.exp_hand(list_exp,newn)==1:
                                    artest=np.array(test)
                                    rnt=np.array(newrt)
                                    ol=np.nonzero(artest==1)
                                    o=np.nonzero(artest==0)
                                    vo=rnt[o]
                                    #raw_input()
                                    for no in range(len(vo)):
                                        #print '%%%%%%%%%'
                                        zi=lawij.mod(vo[no])
                                        #print zi
                                        #print artest
                                        #raw_input()
                                        raxep.append((zi))
                            else:
                                #print f.exp_ync(amp,ampref,testc3)
                                for no in range(len(nt)):
                                    raxep.append(newrt[no])
                                #print 'accepted'
                        else:
                            #print 'ouuuuut'
                            #print f.exp_hand(list_exp,newn)
                            if f.exp_hand(list_exp,newn)==1 :
                                artest=np.array(test)
                                artestc3=np.array(testc3)
                                rnt=np.array(newrt)
                                #print nt
                                #raw_input()
                                ol=np.nonzero(artest==1)
                                condlist=[artest==0,artestc3==0]
                                g=condlist[0]*condlist[1]
                                #print g,artest,artestc3
                                o=np.nonzero(artest==0)
                                #print '%%%hhh%%%%%%'
                                #print rnt
                                #print test
                                vo=rnt[g]
                                #print vo,'%%%hhh%%%%%%'
                                #raw_input()
                                for no in range(len(vo)):
                                            #print '%%%%%%%%%'
                                            zi=lawij.mod(vo[no])
                                            #print zi
                                            #print artest
                                            #raw_input()
                                            raxep.append((zi))
                    else:
                                #print f.exp_ync(amp,ampref,testc3)
                                for no in range(len(nt)):
                                    raxep.append(newrt[no])
                                #print 'accepted'

            else:
                if f.exp_hand(list_exp,newn)==1 :
                    artest=np.array(test)
                    #print newrt
                    rnt=xml_gen.cell2tab8(newrt)
                    #print nt
                    #raw_input()
                    ol=np.nonzero(artest==1)
                    o=np.nonzero(artest==0)
                    #print '%%%hhh%%%%%%'
                    rnt=np.array(rnt)
                    #print test
                    vo=rnt[o]
                    #print vo
                    #print len(vo)
                    #raw_input()
                    for no in range(len(vo)):
                                #print '%%%%%%%%%'
                                zi=lawij.mod(vo[no])
                                #print zi
                                #print artest
                                #raw_input()
                                raxep.append((zi))
                    #print 'may be...'
                #else:
                    #print 'refused without discussion'
            #print test

            #print 'DONE'
            #print '%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%'
            #print '        '
            #raw_input()
    locale.setlocale(locale.LC_ALL, old_locale)
    return raxep,tabc
    #for lll in range(len(raxep)):
        #print  raxep[lll]


from pylivetrader import api
from pylivetrader.assets import Asset

import pylivetrader.protocol as zp
from pylivetrader.assets import Equity
from pylivetrader.finance.order import (
    Order as ZPOrder,
)
from pylivetrader.backend.base import BaseBackend
import pandas as pd
import numpy as np
import string
from trading_calendars import get_calendar

from logbook import Logger

log = Logger(__name__)

MAX_FAKE_BARS = 3000

def _num_to_symbol(n):
    buf = []
    while n >= 0:
        a = n % 26
        b = n // 26 - 1
        c = string.ascii_uppercase[a]
        buf.append(c)
        n = b
    return ''.join(reversed(buf))

def _check_fill(order, price_df):
    price = price_df.close.values[-1]
    if order.amount > 0:
        max_price = price_df.high.max()
        if order.limit is not None and order.stop is not None:
            # STOP LIMIT
            return price <= order.limit and max_price >= order.stop
        elif order.stop is not None:
            # STOP
            return max_price >= order.stop
        elif order.limit is not None:
            # LIMIT
            return price <= order.limit
        else:
            # MARKET
            return True
    else:
        min_price = price_df.low.min()
        if order.limit is not None and order.stop is not None:
            # STOP LIMIT
            return price >= order.limit and min_price <= order.stop
        elif order.stop is not None:
            # STOP
            return min_price <= order.stop
        elif order.limit is not None:
            # LIMIT
            return price >= order.limit
        else:
            # MARKET
            return True
    return False

class Backend(BaseBackend):
    '''This backend is a minimal simulator with
    naive order filling for mainly smoke testing.
    The price data is supplied by fake data generator.
    '''

    def __init__(self, cash=1e6, size=50, clock=None):
        '''
        paramters:
            cash: initial cash balance
            size: the number of stocks in universe
            clock: soft clock
        '''
        self._account = zp.Account()
        self._account.buying_power = cash
        self._account.total_position_value = 0
        self._portfolio = zp.Portfolio()
        self._portfolio.cash = cash
        self._portfolio.positions = self._positions = zp.Positions()
        self._order_seq = 0
        self._orders = {}
        self._clock = clock
        self._last_process_time = None

        self._data_proxy = FakeDataBackend(size=size, clock=clock)

    @property
    def now(self):
        if self._clock is not None:
            return self._clock.now
        return pd.Timestamp.now(tz='America/New_York')

    def _fill(self, order, price):
        '''Fill the order at this price and delete the order
        from the pending list.'''
        del self._orders[order.id]
        pos = self._positions.pop(order.asset, None)
        if pos is None:
            pos = zp.Position(order.asset)
            pos.amount = 0
            pos.cost_basis = 0

        new_amount = pos.amount + order.amount
        if new_amount < 0:
            log.warn('not enough shares')
            return
        volume = price * order.amount
        self._portfolio.cash -= volume
        self._account.cash = self._portfolio.cash
        self._account.buying_power -= volume

        if new_amount == 0:
            return

        new_volume = price * order.amount + pos.cost_basis * pos.amount
        pos.cost_basis = new_volume / new_amount
        pos.amount = new_amount
        pos.last_sale_date = self.now
        self._positions[order.asset] = pos

    def _process_orders(self):
        if self._last_process_time == self.now:
            return
        # materialize the list so that _fill() can delete entry.
        for order in list(self._orders.values()):
            asset = order.asset
            price_df = self.get_bars([asset], '1m')[asset]
            price_df = price_df[order.dt:]

            if _check_fill(order, price_df):
                self._fill(order, price_df.close.values[-1])

        posval = self._portfolio.positions_value = sum([
            p.cost_basis * p.amount for p in self._positions.values()
        ])
        self._portfolio.portfolio_value = self._portfolio.cash + posval
        self._last_process_time = self.now

    def set_position(self, symbol, amount, cost_basis,
                     last_sale_price=None, last_sale_date=None):
        asset = api.symbol(symbol) if isinstance(symbol, str) else symbol
        pos = zp.Position(asset)
        pos.amount = amount
        pos.cost_basis = cost_basis
        pos.last_sale_price = last_sale_price
        pos.last_sale_date = last_sale_date
        self._positions[asset] = pos

    @property
    def positions(self):
        self._process_orders()
        return self._positions

    @property
    def portfolio(self):
        self._process_orders()
        return self._portfolio

    @property
    def account(self):
        return self._account

    @property
    def orders(self):
        self._process_orders()
        return self._orders

    def batch_order(self, args):
        return [self.order(*order) for order in args]

    def order(self, asset, amount, style):
        if amount == 0:
            return
        limit_price = style.get_limit_price(amount > 0) or None
        stop_price = style.get_stop_price(amount > 0) or None
        self._order_seq += 1
        zpOrder = ZPOrder(
            dt=self.get_last_traded_dt(asset),
            asset=asset,
            amount=amount,
            limit=limit_price,
            stop=stop_price,
            id=self._order_seq,
        )
        self._orders[zpOrder.id] = zpOrder
        return zpOrder

    def cancel_order(self, zp_order_id):
        if zp_order_id in self._orders:
            del self._orders[zp_order_id]

    def get_equities(self):
        return self._data_proxy.get_equities()

    def get_last_traded_dt(self, asset):
        return self._data_proxy.get_last_traded_dt(asset)

    def get_spot_value(
            self,
            assets,
            field,
            dt,
            data_frequency,
            quantopian_compatible=True):
        return self._data_proxy.get_spot_value(
            assets, field, dt, data_frequency)

    def get_bars(self, assets, data_frequency, bar_count=500):
        return self._data_proxy.get_bars(assets, data_frequency, bar_count)

class FakeDataBackend:
    '''A data backend that generates synthesic sin wave price data
    for a synthesically generated fixed universe.
    '''

    def __init__(self, size=50, clock=None):
        self._size = size
        self._clock = clock
        self._cal = get_calendar('NYSE')
        self._fake_bars = {
            '1d': {}, '1m': {},
        }

    def _populate_missing(self, assets, data_frequency):
        '''Populates cache bars that are not filled yet.'''
        missing = []
        for asset in assets:
            if asset not in self._fake_bars[data_frequency]:
                missing.append(asset)

        if len(missing) == 0:
            return

        fake_end = self._clock.end_time
        for asset in missing:
            if data_frequency == '1m':
                mask = self._cal.all_minutes
                end = fake_end
            else:
                mask = self._cal.all_sessions
                end = fake_end.floor('1D')
        mask = mask[mask <= end]
        if len(mask) >= MAX_FAKE_BARS:
            mask = mask[-MAX_FAKE_BARS:]
        mask = mask.tz_convert('America/New_York')

        for asset in missing:
            scale = asset.sid
            x = np.linspace(0, len(mask), len(mask))
            ts = np.abs(np.sin(x)) + 1.0
            df = pd.DataFrame({
                'open': ts * scale,
                'high': ts * scale + 0.5,
                'low': ts * scale - 0.5,
                'close': ts * scale + 0.1,
                'volume': (ts * scale * 1e6).astype(int),
            }, index=mask)

            self._fake_bars[data_frequency][asset] = df

    def get_equities(self):
        return [
            Equity(
                sid=i + 1,
                symbol=_num_to_symbol(i),
                asset_name='Test {}'.format(_num_to_symbol(i)),
                exchange='NYSE',
                start_date=pd.Timestamp('1970-01-01', tz='utc'),
                end_date=pd.Timestamp('2050-01-01', tz='utc'),
            ) for i in range(self._size)
        ]

    def get_last_traded_dt(self, asset):
        return self.now

    def get_spot_value(
            self,
            assets,
            field,
            dt,
            data_frequency,
            quantopian_compatible=True):
        now = self.now

        def _get_for_symbol(df, field):
            _field = field
            if field == 'last_traded':
                return df.index[-1]
            elif field == 'price':
                _field = 'close'
            return df[_field].values[-1]

        if isinstance(assets, Asset):
            df = self.get_bars([assets], '1m')[assets][:now]
            return _get_for_symbol(df, field)

        dfs = self.get_bars(assets, '1m')
        return [
            _get_for_symbol(dfs[asset][:now], field) for asset in assets
        ]

    def get_bars(self, assets, data_frequency, bar_count=500):
        now = self.now

        self._populate_missing(assets, data_frequency)

        items = []
        for asset in assets:
            df = self._fake_bars[data_frequency][asset]
            df = df[df.index <= now]
            df = df.iloc[-bar_count:]
            df.columns = pd.MultiIndex.from_product([[asset, ], df.columns])
            items.append(df)
        return pd.concat(items, axis=1)

    @property
    def now(self):
        if self._clock is not None:
            return self._clock.now
        return pd.Timestamp.now(tz='America/New_York')


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import time

import numpy as np
import tensorflow as tf
import os
import random
import copy

from tensorflow.python.util import nest
from discriminator import *
from config import *

tf.app.flags.DEFINE_string('data_path', "./data",
                           """Path where the data will be loaded.""")
tf.app.flags.DEFINE_string('name', "mysubmission",
                           """Path where the data will be loaded.""")
tf.app.flags.DEFINE_string('model_architecture', 'mlp_1_img_1_512_0',
                           """Number of images to process in a batch.""")
tf.app.flags.DEFINE_integer('epochs', 30,
                           """Number of epochs.""")
args = tf.app.flags.FLAGS

def data_loader(data_path=None, data_type = '_full', use_mc_samples=False):
    """
    Data format (compatible with Show Attend and Tell):
    the data file is a dict has the following keys:
    'file_names'
    'image_idxs'
    'captions': a dict has keys 'gen' for generator and 'dis' for discriminator
    'features': a dict has keys 'gen' for generator and 'dis' for discriminator
                (to be loaded when needed)
    'word_to_idx': a dict with word to idx mapping
    """
    data_train = np.load(os.path.join(data_path, "data_train_full.npy")).item()
    data_val = np.load(os.path.join(data_path, "data_val_full.npy")).item()
    data_test = np.load(os.path.join(data_path, "data_test_full.npy")).item()
    if use_mc_samples:
        # mc_train = np.load(os.path.join(data_path, 'dumped_train.npy')).item()
        # mc_val = np.load(os.path.join(data_path, 'dumped_val.npy')).item()
        # mc_test = np.load(os.path.join(data_path, 'dumped_test.npy')).item()
        mc_train = np.load('./data/dumped_train.npy').item()
        mc_val = np.load('./data/dumped_val.npy').item()
        mc_test = np.load('./data/dumped_test.npy').item()
        data_train = add_mc_samples(data_train, mc_train)
        data_val = add_mc_samples(data_val, mc_val)
        data_test = add_mc_samples(data_test, mc_test)

    data_train['features']['dis'] = np.load(
            './data/resnet152/feature_dis_train%s.npy' % (data_type)
        ).item()
    data_val['features']['dis'] = np.load(
            './data/resnet152/feature_dis_val%s.npy' % (data_type)
        ).item()
    data_test['features']['dis'] = np.load(
            './data/resnet152/feature_dis_test%s.npy' % (data_type)
        ).item()

    word_embedding = np.load(
            './data/word_embedding_%s.npy' % (str(Config().embedding_size))
        )
    return [data_train, data_val, data_test, word_embedding]

def main(_):
    exp_name = "%s_scoring"%(args.name)
    log_path = './log/' + exp_name
    save_path = './model/' + exp_name
    data_path = os.path.join(args.data_path, args.name)
    if not os.path.exists(log_path):
        os.makedirs(log_path)
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    train_models = [args.name]
    test_models = [args.name, 'human']

    [data_train, data_val, data_test, word_embedding] = data_loader(
            data_path, use_mc_samples=True)
    word_to_idx = data_train['word_to_idx']

    config = Config()
    config = config_model_coco(config, args.model_architecture)
    config.max_epoch = args.epochs

    print("Model architecture:%s"%(args.model_architecture))
    with tf.Graph().as_default():
        with tf.name_scope("Train"):
            with tf.variable_scope("Discriminator", reuse=None):
                mtrain = Discriminator(word_embedding, word_to_idx, use_glove=True,
                                       config=config, is_training=True)
            tf.summary.scalar("Training Loss", mtrain._loss)
            tf.summary.scalar("Training Accuracy", mtrain._accuracy)

        with tf.name_scope("Val"):
            with tf.variable_scope("Discriminator", reuse=True):
                mval = Discriminator(word_embedding, word_to_idx, use_glove=True,
                                     config=config, is_training=False)
            tf.summary.scalar("Validation Loss", mval._loss)
            tf.summary.scalar("Validation Accuracy", mval._accuracy)

        config_sess = tf.ConfigProto(allow_soft_placement=True)
        config_sess.gpu_options.allow_growth = True
        with tf.Session(config=config_sess) as sess:
            tf.global_variables_initializer().run()
            summary_writer = tf.summary.FileWriter(log_path, graph=tf.get_default_graph())
            saver = tf.train.Saver()

            output_filename = '%s.txt' % (args.model_architecture)
            output_filepath = os.path.join(save_path, output_filename)
            f = open(output_filepath, 'w')

            # Training
            for i in xrange(config.max_epoch):
                print("Epoch: %d" % (i + 1))
                train_loss, train_acc = train(sess, mtrain, data_train,
                                              gen_model=train_models, epoch=i,
                                              config=config)

                for test_model in test_models:
                    [acc, logits, scores] = inference(
                            sess, mval, data_val, test_model, config=config)
                    s = np.mean(scores[:,:,0])
                    f.write("%f\t" % s)
                    a = np.mean(acc)
                    f.write("%f\t" % a)
                f.write("\n")
            f.close()

            if save_path:
                model_path = os.path.join(save_path, args.model_architecture)
                print("Saving model to %s." % model_path)
                saver.save(sess, model_path, global_step=i+1)
                print("Model saved to %s." % model_path)

if __name__ == "__main__":
    tf.app.run()