"""

Author: Gurkirt Singh
Modified from https://github.com/gurkirt/realtime-action-detection
Licensed under The MIT License [see LICENSE for details]

"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.init as init
import argparse
from torch.autograd import Variable
import torch.utils.data as data
from data import v2, ActionDetection, NormliseBoxes, detection_collate, CLASSES, BaseTransform
from utils.augmentations import SSDAugmentation
from layers.modules import MultiboxLoss
from layers.functions import PriorBox
from layers import MatchPrior
from AMTNet import AMTNet
import numpy as np
import time, pdb
from utils.evaluation import evaluate_detections
from layers.box_utils import nms, decode_seq
from utils import  AverageMeter
from torch.optim.lr_scheduler import MultiStepLR
# from torchviz import make_dot, make_dot_from_trace

def str2bool(v):
    return v.lower() in ("yes", "true", "t", "1")

parser = argparse.ArgumentParser(description='AMTNet detection training script')
parser.add_argument('--version', default='v2', help='conv11_2(v2) or pool6(v1) as last layer')
parser.add_argument('--basenet', default='vgg16_reducedfc.pth', help='pretrained base model')
parser.add_argument('--dataset', default='ucf24', help='pretrained base model')
parser.add_argument('--train_split', default=1, type=int, help='Split id')
parser.add_argument('--ssd_dim', default=300, type=int, help='Input Size for SSD') # only support 300 now
parser.add_argument('--seq_len', default=2, type=int, help='Input sequence length ')
parser.add_argument('--seq_gap', default=0, type=int, help='Gap between the frame of sequence')
parser.add_argument('--fusion_type', default='cat', type=str,
                    help='Fusion type to fuse from sequence of frames; options are SUM, CAT and NONE')
                    #
parser.add_argument('--input_type_base', default='rgb', type=str, help='INput tyep default rgb can take flow (brox or fastOF) as well')
parser.add_argument('--input_type_extra', default='brox', type=str, help='INput tyep default brox can take flow (brox or fastOF) as well')
parser.add_argument('--input_frames_base', default=1, type=int, help='Number of input frame, default for rgb is 1')
parser.add_argument('--input_frames_extra', default=5, type=int, help='Number of input frame, default for flow is 5')
parser.add_argument('--jaccard_threshold', default=0.5, type=float, help='Min Jaccard index for matching')
parser.add_argument('--batch_size', default=4, type=int, help='Batch size for training')
parser.add_argument('--num_workers','-j', default=4, type=int, help='Number of workers used in dataloading')
parser.add_argument('--max_iter', default=40000, type=int, help='Number of training iterations')
parser.add_argument('--val_step', default=10000, type=int, help='Number of training iterations')
parser.add_argument('--cuda', default=1, type=str2bool, help='Use cuda to train model')
parser.add_argument('--ngpu', default=1, type=int, help='Use cuda to train model')
parser.add_argument('--lr', '--learning-rate', default=0.0005, type=float, help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
parser.add_argument('--stepvalues', default='10000,30000', type=str, help='step points for learning rate drop')
parser.add_argument('--weight_decay', default=1e-4, type=float, help='Weight decay for SGD')
parser.add_argument('--gamma', default=0.1, type=float, help='Gamma update for SGD at for stepwise schedule')
parser.add_argument('--visdom', default=False, type=str2bool, help='Use visdom to for loss visualization')
parser.add_argument('--vis_port', default=8095, type=int, help='Port for Visdom Server')
parser.add_argument('--data_root', default='~/data/', help='Location of where in data is located like images and annotation file')
parser.add_argument('--save_root', default='~/cache/', help='Location to where we wanr save the checkpoints of models')
parser.add_argument('--iou_thresh', default=0.5, type=float, help='Evaluation threshold')
parser.add_argument('--conf_thresh', default=0.01, type=float, help='Confidence threshold for evaluation')
parser.add_argument('--nms_thresh', default=0.45, type=float, help='NMS threshold')
parser.add_argument('--default_mult', default=1.0, type=float, help='NMS threshold')
parser.add_argument('--topk', default=50, type=int, help='topk for evaluation')
parser.add_argument('--man_seed', default=123, type=int, help='manula seed')
args = parser.parse_args()

import socket
import getpass
username = getpass.getuser()
hostname = socket.gethostname()

print('\n\n ', username, ' is using ', hostname, '\n\n')
if hostname == 'mars':
    args.data_root = '/mnt/mars-fast/datasets/'
    args.save_root = '/mnt/mars-gamma/'
    args.vis_port = 8097
elif hostname in ['sun']:
    args.data_root = '/mnt/sun-gamma/'
    args.save_root = '/mnt/sun-gamma/'
    args.vis_port = 8096
elif hostname == 'mercury':
    args.data_root = '/mnt/mercury-fast/datasets/'
    args.save_root = '/mnt/mercury-beta/'
    args.vis_port = 8098
elif username == 'gurkirt' and hostname.startswith('comp'):
    args.data_root = '/home/gurkirt/datasets/'
    args.save_root = '/home/gurkirt/cache/'
    args.vis_port = 8097
    visdom=False
# python train.py --seq_len=2 --num_workers=4 --batch_size=16 --ngpu=2 --fusion_type=NONE --input_type_base=brox --input_frames_base=5 --stepvalues=30000,50000 --max_iter=60000 --val_step=10000 --lr=0.001

torch.set_default_tensor_type('torch.FloatTensor')
np.random.seed(args.man_seed)
torch.manual_seed(args.man_seed)
if args.cuda:
    torch.cuda.manual_seed_all(args.man_seed)

def print_node(gdf):
    node_fns = gdf.next_functions
    for fn in node_fns:
        print(fn)
        print_node(fn[0][0])

def main():
    args.cfg = v2
    args.train_sets = 'train'
    args.test_sets = 'test'
    kd = 3
    args.means = (104, 117, 123)
    num_classes = len(CLASSES[args.dataset]) + 1 # only support multiclass datasets, not multilabel
    args.num_classes = num_classes
    args.stepvalues = [int(val) for val in args.stepvalues.split(',')]
    args.loss_reset_step = 30
    # args.val_step = 30000
    args.print_step = 10
    args.fusion_type = args.fusion_type.lower()
    args.fusion = args.fusion_type in ['sum','cat','mean']
    ## Define the experiment Name will used for save directory and ENV for visdom
    if not args.fusion:
        args.exp_name = 'AMTNet-{}-s{:d}-{}-sl{:02d}sg{:02d}-bs{:02d}-lr{:05d}'.format(args.dataset, args.train_split,
                                                                                args.input_type_base,
                                                                                args.seq_len, args.seq_gap,
                                                                                args.batch_size, int(args.lr * 100000))
    else:
        args.exp_name = 'AMTNet-{}-s{:d}-{}-{}-{}-sl{:02d}sg{:02d}-bs{:02d}-lr{:05d}'.format(args.dataset, args.train_split,
                                                                                args.fusion_type, args.input_type_base,
                                                                                args.input_type_extra,
                                                                                args.seq_len, args.seq_gap,
                                                                                args.batch_size,int(args.lr * 100000))

    num_feat_multiplier = {'cat': 2, 'sum': 1, 'mean': 1, 'none': 1}
    # fusion type can one of the above keys
    args.fmd = [512, 1024, 512, 256, 256, 256]
    args.kd = 3
    args.fusion_num_muliplier = num_feat_multiplier[args.fusion_type]

    ## DEFINE THE NETWORK
    net = AMTNet(args)

    if args.fusion:
        base_weights = torch.load(args.save_root +'weights/AMTNet_single_stream_{:s}_s{:02d}.pth'.format(args.input_type_base, args.train_split))
        extra_weights = torch.load(args.save_root + '/weights/AMTNet_single_stream_{:s}_s{:02d}.pth'.format(args.input_type_extra, args.train_split))
        print('Loading base network...')
        net.core_base.load_my_state_dict(base_weights, input_frames=args.input_frames_base)
        net.core_extra.load_my_state_dict(extra_weights, input_frames=args.input_frames_extra)
    else:
        base_weights = torch.load(args.data_root +'/weights/{}-ssd300_ucf24_120000.pth'.format(args.input_type_base))
        net.core_base.load_my_state_dict(base_weights, input_frames=args.input_frames_base)

    args.data_root += args.dataset + '/'
    args.save_root += args.dataset + '/'

    net = net.cuda()

    def xavier(param):
        init.xavier_uniform_(param)

    def weights_init(m):
        if isinstance(m, nn.Conv2d):
            xavier(m.weight.data)
            m.bias.data.zero_()
        elif isinstance(m, nn.Linear):
            xavier(m.weight.data)
            m.bias.data.zero_()

    print('Initializing weights for HEADs...')
    net.loc.apply(weights_init)
    net.conf.apply(weights_init)

    args.save_root = args.save_root + 'cache/' + args.exp_name + '/'
    if not os.path.isdir(args.save_root):
        os.makedirs(args.save_root)

    if args.ngpu>1:
        print('\nLets do dataparallel\n\n')
        net = torch.nn.DataParallel(net)

    parameter_dict = dict(net.named_parameters()) # Get parmeter of network in dictionary format wtih name being key
    params = []

    #Set different learning rate to bias layers and set their weight_decay to 0
    mult = 1; decay = 0

    for name, param in parameter_dict.items():
        if name.find('bias') > -1:
            mult = 2.0; decay = 0
        else:
            mult = 1.0;  decay = 1
        if name.find('vgg')> -1 or name.find('extra')>-1 or name.find('L2Norm')>-1:
            mult = mult/args.seq_len

        # print(name, 'layer parameters will be trained @ {}'.format(args.lr*mult))
        params += [{'params':[param], 'lr': args.lr*mult, 'weight_decay':args.weight_decay*decay}]

    optimizer = optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    criterion = MultiboxLoss()

    scheduler = MultiStepLR(optimizer, milestones=args.stepvalues, gamma=args.gamma)
    # Get proior or anchor boxes
    with torch.no_grad():
        priorbox = PriorBox(v2, args.seq_len)
        priors = priorbox.forward()
    train(args, net, priors, optimizer, criterion, scheduler)

def train(args, net, priors, optimizer, criterion, scheduler):
    log_file = open(args.save_root+"training.log", "w", 1)
    log_file.write(args.exp_name+'\n')
    for arg in sorted(vars(args)):
        print(arg, getattr(args, arg))
        log_file.write(str(arg)+': '+str(getattr(args, arg))+'\n')

    net.train()
    # loss counters
    batch_time = AverageMeter()
    losses = AverageMeter()
    loc_losses = AverageMeter()
    cls_losses = AverageMeter()

    print('Loading Dataset...')
    train_dataset = ActionDetection(args, args.train_sets, SSDAugmentation(args.ssd_dim, args.means),
                                    NormliseBoxes(), anno_transform=MatchPrior(priors, args.cfg['variance']))
    log_file.write(train_dataset.print_str)
    print(train_dataset.print_str)
    val_dataset = ActionDetection(args, args.test_sets, BaseTransform(args.ssd_dim, args.means),
                                  NormliseBoxes(), full_test=False)
    log_file.write(val_dataset.print_str)
    # print(val_dataset.print_str)
    epoch_size = len(train_dataset) // args.batch_size

    print('Training SSD on', train_dataset.name)

    if args.visdom:

        import visdom
        viz = visdom.Visdom(env=args.exp_name, port=args.vis_port)
        # initialize visdom loss plot
        lot = viz.line(
            X=torch.zeros((1,)).cpu(),
            Y=torch.zeros((1, 6)).cpu(),
            opts=dict(
                xlabel='Iteration',
                ylabel='Loss',
                title='Current SSD Training Loss',
                legend=['REG', 'CLS', 'AVG', 'S-REG', ' S-CLS', ' S-AVG']
            )
        )
        # initialize visdom meanAP and class APs plot
        legends = ['meanAP']
        for cls in CLASSES[args.dataset]:
            legends.append(cls)
        val_lot = viz.line(
            X=torch.zeros((1,)).cpu(),
            Y=torch.zeros((1,args.num_classes)).cpu(),
            opts=dict(
                xlabel='Iteration',
                ylabel='Mean AP',
                title='Current SSD Validation mean AP',
                legend=legends
            )
        )

    batch_iterator = None
    train_data_loader = data.DataLoader(train_dataset, args.batch_size, num_workers=args.num_workers,
                                  shuffle=True, collate_fn=detection_collate, pin_memory=True)
    val_data_loader = data.DataLoader(val_dataset, args.batch_size, num_workers=args.num_workers,
                                 shuffle=False, collate_fn=detection_collate, pin_memory=True)
    itr_count = 0
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for iteration in range(args.max_iter + 1):
        if (not batch_iterator) or (iteration % epoch_size == 0):
            # create batch iterator
            batch_iterator = iter(train_data_loader)

        # load train data
        images, _ , prior_gt_labels, prior_gt_locations, _, _ = next(batch_iterator)
        # images, ground_truths, _ , _, num_mt, img_indexs
        # pdb.set_trace()
        images = [img.cuda(0, non_blocking=True) for img in images if not isinstance(img, list)]
        prior_gt_labels = prior_gt_labels.cuda(0, non_blocking=True)
        prior_gt_locations = prior_gt_locations.cuda(0, non_blocking=True)
        # forward
        cls_out, reg_out = net(images)

        optimizer.zero_grad()
        loss_l, loss_c = criterion(cls_out, reg_out, prior_gt_labels, prior_gt_locations)
        loss = loss_l + loss_c

        loss.backward()
        optimizer.step()
        scheduler.step()

        # pdb.set_trace()
        loc_loss = loss_l.item()
        conf_loss = loss_c.item()
        # print('Loss data type ',type(loc_loss))
        loc_losses.update(loc_loss)
        cls_losses.update(conf_loss)
        losses.update((loc_loss + conf_loss)/2.0)

        if iteration == 103:
            loc_losses.reset()
            cls_losses.reset()
            losses.reset()
            batch_time.reset()

        if iteration % args.print_step == 0:
            if args.visdom and iteration>100:
                losses_list = [loc_losses.val, cls_losses.val, losses.val, loc_losses.avg, cls_losses.avg, losses.avg]
                viz.line(X=torch.ones((1, 6)).cpu() * iteration,
                    Y=torch.from_numpy(np.asarray(losses_list)).unsqueeze(0).cpu(),
                    win=lot,
                    update='append')

            torch.cuda.synchronize()
            t1 = time.perf_counter()
            batch_time.update(t1 - t0)

            print_line = 'Itration {:02d}/{:06d}/{:06d} loc-loss {:.3f}({:.3f}) cls-loss {:.3f}({:.3f}) ' \
                         'average-loss {:.3f}({:.3f}) Timer {:0.3f}({:0.3f})'.format(iteration//epoch_size,
                          iteration, args.max_iter, loc_losses.val, loc_losses.avg, cls_losses.val,
                          cls_losses.avg, losses.val, losses.avg, batch_time.val, batch_time.avg)

            torch.cuda.synchronize()
            t0 = time.perf_counter()
            log_file.write(print_line+'\n')
            print(print_line)

            itr_count += 1

            if itr_count % args.loss_reset_step == 0 and itr_count > 0:
                loc_losses.reset()
                cls_losses.reset()
                losses.reset()
                batch_time.reset()
                print('Reset ', args.exp_name,' after', itr_count*args.print_step)
                itr_count = 0

        if (iteration % args.val_step == 0 or iteration in [1000, args.max_iter]) and iteration>0:
            torch.cuda.synchronize()
            tvs = time.perf_counter()
            print('Saving state, iter:', iteration)
            torch.save(net.state_dict(), args.save_root + 'AMTNet_' +
                       repr(iteration) + '.pth')

            net.eval() # switch net to evaluation mode
            with torch.no_grad():
                mAP, ap_all, ap_strs = validate(args, net, priors, val_data_loader, val_dataset, iteration, iou_thresh=args.iou_thresh)

            for ap_str in ap_strs:
                print(ap_str)
                log_file.write(ap_str+'\n')
            ptr_str = '\nMEANAP:::=>'+str(mAP)+'\n'
            print(ptr_str)
            log_file.write(ptr_str)

            if args.visdom:
                aps = [mAP]
                for ap in ap_all:
                    aps.append(ap)
                viz.line(
                    X=torch.ones((1, args.num_classes)).cpu() * iteration,
                    Y=torch.from_numpy(np.asarray(aps)).unsqueeze(0).cpu(),
                    win=val_lot,
                    update='append'
                        )
            net.train() # Switch net back to training mode
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            prt_str = '\nValidation TIME::: {:0.3f}\n\n'.format(t0-tvs)
            print(prt_str)
            log_file.write(ptr_str)

    log_file.close()

def validate(args, net, priors, val_data_loader, val_dataset, iteration_num, iou_thresh=0.5):
    """Test a SSD network on an image database."""
    print('Validating at ', iteration_num)
    num_images = len(val_dataset)
    num_classes = args.num_classes
    priors = priors.cuda()
    det_boxes = [[] for _ in range(len(CLASSES[args.dataset]))]
    gt_boxes = []
    print_time = True
    batch_iterator = None
    val_step = 100
    count = 0
    torch.cuda.synchronize()
    ts = time.perf_counter()
    softmax = nn.Softmax(dim=2).cuda()
    with torch.no_grad():
        for val_itr in range(len(val_data_loader)):
            if not batch_iterator:
                batch_iterator = iter(val_data_loader)

            torch.cuda.synchronize()
            t1 = time.perf_counter()

            images, ground_truths, _ , _, num_mt, img_indexs = next(batch_iterator)
            batch_size = images[0].size(0)
            #images = images.permute(1, 0, 2, 3, 4)
            height, width = images[0].size(3), images[0].size(4)

            images = [img.cuda(0, non_blocking=True) for img in images if not isinstance(img, list)]

            conf_preds, loc_data = net(images)

            # pdb.set_trace()
            conf_scores_all = softmax(conf_preds).clone()

            if print_time and val_itr%val_step == 0:
                torch.cuda.synchronize()
                tf = time.perf_counter()
                print('Forward Time {:0.3f}'.format(tf-t1))

            for b in range(batch_size):
                # pdb.set_trace()
                inds = np.asarray([m*args.seq_len for m in range(num_mt[b])])
                # pdb.set_trace()
                gt = ground_truths[b].numpy()
                gt = gt[inds]
                gt[:,0] *= width
                gt[:,2] *= width
                gt[:,1] *= height
                gt[:,3] *= height
                gt_boxes.append(gt)
                decoded_boxes = decode_seq(loc_data[b], priors, args.cfg['variance'], args.seq_len)
                decoded_boxes = decoded_boxes[:,:4].clone()
                conf_scores = conf_scores_all[b].clone()
                #Apply nms per class and obtain the results
                for cl_ind in range(1, num_classes):
                    # pdb.set_trace()
                    scores = conf_scores[:, cl_ind].squeeze()
                    c_mask = scores.gt(args.conf_thresh)  # greater than minmum threshold
                    scores = scores[c_mask].squeeze() # reduce the dimension so if no element then # of dim is 0
                    if scores.dim() == 0:
                        det_boxes[cl_ind - 1].append(np.asarray([]))
                        continue
                    boxes = decoded_boxes.clone()
                    l_mask = c_mask.unsqueeze(1).expand_as(boxes)
                    boxes = boxes[l_mask].view(-1, 4)
                    # idx of highest scoring and non-overlapping boxes per class
                    ids, counts = nms(boxes, scores, args.nms_thresh, args.topk)  # idsn - ids after nms
                    scores = scores[ids[:counts]].cpu().numpy()
                    boxes = boxes[ids[:counts]].cpu().numpy()
                    # print('boxes sahpe',boxes.shape)
                    boxes[:,0] *= width
                    boxes[:,2] *= width
                    boxes[:,1] *= height
                    boxes[:,3] *= height

                    for ik in range(boxes.shape[0]):
                        boxes[ik, 0] = max(0, boxes[ik, 0])
                        boxes[ik, 2] = min(width, boxes[ik, 2])
                        boxes[ik, 1] = max(0, boxes[ik, 1])
                        boxes[ik, 3] = min(height, boxes[ik, 3])
                    cls_dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=True)
                    det_boxes[cl_ind-1].append(cls_dets)
                count += 1

            if val_itr%val_step == 0:
                torch.cuda.synchronize()
                te = time.perf_counter()
                print('im_detect: {:d}/{:d} time taken {:0.3f}'.format(count, num_images, te-ts))
                torch.cuda.synchronize()
                ts = time.perf_counter()
            if print_time and val_itr%val_step == 0:
                torch.cuda.synchronize()
                te = time.perf_counter()
                print('NMS stuff Time {:0.3f}'.format(te - tf))
    print('Evaluating detections for itration number ', iteration_num)
    return evaluate_detections(gt_boxes, det_boxes, CLASSES[args.dataset], iou_thresh=iou_thresh)

if __name__ == '__main__':
    main()


import functools
import json
import subprocess
import threading
import time

import numpy as np
import typing
from PIL import Image

from .basic_predictor import BasicPredictor

class GazeMLPredictor(BasicPredictor):
    """GazeML based predictor

    Actually, it is just an interface for gazeML subprocess as i was too lazy to implement
    in a proper way.
    All the information is got from gazeml subprocess.

    :ivar gazeML_process: (subprocess.Popen), running gazeml subprocess
    :ivar eye_one_target: (np.ndarray, shape [3]), gaze vector from first eye
    :ivar eye_two_target: (np.ndarray, shape [3]), gaze vector from second eye
    :ivar np_points: (np.ndarray, shape [68, 2]), current facial landmarks
    """
    def __init__(self):
        """Constructor

        Runs the subprocess and a thread for reading the data from subprocess
        """
        self.gazeML_process = subprocess.Popen(
            ["python", "./from_internet_or_for_from_internet/GazeML/src/elg_demo.py"],
            stdout=subprocess.PIPE)

        self.eye_one_target = np.array([0, 0, 1])
        self.eye_two_target = np.array([0, 0, 1])
        self.np_points = np.zeros((68, 2))

        def read_from_subprocess(predictor: GazeMLPredictor, process: subprocess.Popen):
            """Read info from subprocess stdout and change the predictor values according to it

            :param predictor: predictor to write
            :param process: process where to read from
            :return:
            """
            while True:
                line = process.stdout.readline()
                try:
                    data = json.loads(line)

                    # eye number 0 is second eye
                    # it was observed that y should be inverted
                    if data["eye_number"] == 0:
                        predictor.eye_two_target = np.array([data["x"], -data["y"], data["z"]])
                    else:
                        predictor.eye_one_target = np.array([data["x"], -data["y"], data["z"]])

                    predictor.np_points = np.array(data['np_points']).reshape((68, 2))
                except json.decoder.JSONDecodeError:
                    pass

        thread_func = functools.partial(read_from_subprocess, self, self.gazeML_process)
        thread = threading.Thread(target=thread_func)
        thread.start()

    def predict_eye_vector_and_face_points(
            self,
            imgs: typing.List[Image.Image],
            time_now: float) -> typing.Tuple[
                typing.List[Image.Image],
                typing.List[np.ndarray],
                typing.List[np.ndarray],
                typing.List[np.ndarray],
                dict]:
        """Retrieves the gaze vectors and facial landmarks

        imgs and time_now aren't really needed, they are just for LSP

        :param imgs: list of images to predict from
        :param time_now: current time
        :return: tuple of shape [5], where:
                    [0] - list of resulting images
                    [1] - list of right eye gazes of shape [3]
                    [2] - list of left eye gazes of shape [3]
                    [3] - list of facial landmark points [68,2]
                    [4] - dict of some output information
        """
        time.sleep(0.2)
        return [Image.fromarray(np.zeros((720, 1280)))], [self.eye_one_target],\
               [self.eye_two_target], [self.np_points], {}


from absl import app, flags, logging
from absl.flags import FLAGS

import tensorflow as tf
import numpy as np
import cv2
from tensorflow.keras.callbacks import (
    ReduceLROnPlateau,
    EarlyStopping,
    ModelCheckpoint,
)
from yolov3.models import (
    YoloV3, YoloV3Tiny, YoloLoss,
    yolo_anchors, yolo_anchor_masks,
    yolo_tiny_anchors, yolo_tiny_anchor_masks
)
from yolov3.utils import freeze_all
import yolov3.dataset as dataset

flags.DEFINE_string('train_images', '', 'path to training set')
flags.DEFINE_string('train_labels', '', 'path to training set label file')
flags.DEFINE_string('val_images', '', 'path to validation set')
flags.DEFINE_string('val_labels', '', 'path to validation set label file')
flags.DEFINE_string('classes', './data/face.names', 'path to classes file')
flags.DEFINE_integer('size', 416, 'image size')
flags.DEFINE_boolean('tiny', False, 'yolov3 or yolov3-tiny')
flags.DEFINE_enum('mode', 'fit', ['fit', 'eager'],
                  'fit: model.fit, '
                  'eager: custom GradientTape')
flags.DEFINE_integer('epochs', 50, 'number of epochs')
flags.DEFINE_integer('batch_size', 16, 'batch size')
flags.DEFINE_float('learning_rate', 1e-3, 'learning rate')
flags.DEFINE_enum('freeze', 'none', ['none', 'darknet', 'no_output'],
                  'Freeze none, '
                  'Freeze darknet layers, '
                  'Freeze all but output layer')
flags.DEFINE_string('weights', None,
                    'path to weights file')
flags.DEFINE_integer('weights_num_classes', None, 'specify num class for `weights` file'
                     'when do transfer learning with different number of classes task, '
                     'e.g. darknet weights_num_classes = 80')
flags.DEFINE_enum('transfer', 'no_output', ['no_output', 'darknet'],
                  'Transfer darknet, '
                  'Transfer all but output')

def main(_argv):
    num_classes = len([c.strip() for c in open(FLAGS.classes).readlines()])
    logging.info('classes loaded, number of classes = %d' % num_classes)

    if FLAGS.tiny:
        model = YoloV3Tiny(FLAGS.size, training=True,
                           classes=num_classes)
        anchors = yolo_tiny_anchors
        anchor_masks = yolo_tiny_anchor_masks
    else:
        model = YoloV3(FLAGS.size, training=True, classes=num_classes)
        anchors = yolo_anchors
        anchor_masks = yolo_anchor_masks

    # TensorFlow 2.0/2.1 has a bug when putting finite dateset to model.fit().
    # Walk around by using dataset.repeat() and give the number of steps to model fit.
    # Refer to https://github.com/tensorflow/tensorflow/issues/31509
    training_set_size = dataset.get_dataset_size(FLAGS.train_labels)
    steps_for_train = training_set_size // FLAGS.batch_size
    steps_for_val = dataset.get_dataset_size(FLAGS.val_labels) // FLAGS.batch_size
    logging.info('Training in %d steps, validation in %d steps' % (steps_for_train, steps_for_val))

    if not (FLAGS.train_images and FLAGS.train_labels):
      logging.error('Training images and labels must be specified.')
      return
    train_dataset = dataset.load_yolo_dataset(
        FLAGS.train_images, FLAGS.train_labels, FLAGS.size, shuffle=True, shuffle_buffer_size=training_set_size)
    train_dataset = train_dataset.repeat()
    train_dataset = train_dataset.batch(FLAGS.batch_size, drop_remainder=True)
    train_dataset = train_dataset.map(lambda x, y: (
        dataset.transform_images(x, FLAGS.size),
        dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size)))
    train_dataset = train_dataset.prefetch(
        buffer_size=tf.data.experimental.AUTOTUNE)

    if not (FLAGS.val_images and FLAGS.val_labels):
      logging.error('Validation images and labels must be specified.')
      return
    val_dataset = dataset.load_yolo_dataset(
        FLAGS.val_images, FLAGS.val_labels, FLAGS.size)
    val_dataset = val_dataset.repeat()
    val_dataset = val_dataset.batch(FLAGS.batch_size, drop_remainder=True)
    val_dataset = val_dataset.map(lambda x, y: (
        dataset.transform_images(x, FLAGS.size),
        dataset.transform_targets(y, anchors, anchor_masks, FLAGS.size)))
    val_dataset = train_dataset.prefetch(
        buffer_size=tf.data.experimental.AUTOTUNE)

    # Configure the model
    # Weights configuration
    if FLAGS.weights_num_classes is None:
        if FLAGS.weights:
            model.load_weights(FLAGS.weights)
    else:
        # Transfer learning with incompatible number of classes
        assert FLAGS.weights

        # reset top layers
        if FLAGS.tiny:
            model_pretrained = YoloV3Tiny(
                FLAGS.size, training=True, classes=FLAGS.weights_num_classes or num_classes)
        else:
            model_pretrained = YoloV3(
                FLAGS.size, training=True, classes=FLAGS.weights_num_classes or num_classes)
        model_pretrained.load_weights(FLAGS.weights)

        if FLAGS.transfer == 'darknet':
            model.get_layer('yolo_darknet').set_weights(
                model_pretrained.get_layer('yolo_darknet').get_weights())
        elif FLAGS.transfer == 'no_output':
            for l in model.layers:
                if not l.name.startswith('yolo_output'):
                    l.set_weights(model_pretrained.get_layer(
                        l.name).get_weights())

    # Freeze layers
    if FLAGS.freeze == 'none':
        pass
    elif FLAGS.freeze == 'darknet':
        freeze_all(model.get_layer('yolo_darknet'))
    elif FLAGS.freeze == 'no_output':
        # freeze all but output layers
        for l in model.layers:
            if not l.name.startswith('yolo_output'):
                freeze_all(l)
    model.summary()

    # Configure training process
    optimizer = tf.keras.optimizers.Adam(lr=FLAGS.learning_rate)
    loss = [YoloLoss(anchors[mask], classes=num_classes, ignore_thresh=0.5)
            for mask in anchor_masks]

    if FLAGS.mode == 'eager':
        # Eager mode is great for debugging
        # Non eager graph mode is recommended for real training
        avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32)
        avg_val_loss = tf.keras.metrics.Mean('val_loss', dtype=tf.float32)

        for epoch in range(1, FLAGS.epochs + 1):
            batch = 0
            for images, labels in train_dataset:
                with tf.GradientTape() as tape:
                    outputs = model(images, training=True)
                    regularization_loss = tf.reduce_sum(model.losses)
                    pred_loss = []
                    for output, label, loss_fn in zip(outputs, labels, loss):
                        pred_loss.append(loss_fn(label, output))
                    total_loss = tf.reduce_sum(pred_loss) + regularization_loss

                grads = tape.gradient(total_loss, model.trainable_variables)
                optimizer.apply_gradients(
                    zip(grads, model.trainable_variables))

                logging.info("{}_train_{}, {}, {}".format(
                    epoch, batch, total_loss.numpy(),
                    list(map(lambda x: np.sum(x.numpy()), pred_loss))))
                avg_loss.update_state(total_loss)

                batch += 1
                if batch == steps_for_train:
                    break

            batch = 0
            for images, labels in val_dataset:
                outputs = model(images)
                regularization_loss = tf.reduce_sum(model.losses)
                pred_loss = []
                for output, label, loss_fn in zip(outputs, labels, loss):
                    pred_loss.append(loss_fn(label, output))
                total_loss = tf.reduce_sum(pred_loss) + regularization_loss

                logging.info("{}_val_{}, {}, {}".format(
                    epoch, batch, total_loss.numpy(),
                    list(map(lambda x: np.sum(x.numpy()), pred_loss))))
                avg_val_loss.update_state(total_loss)

                batch += 1
                if batch == steps_for_val:
                    break

            logging.info("{}, train: {}, val: {}".format(
                epoch,
                avg_loss.result().numpy(),
                avg_val_loss.result().numpy()))

            avg_loss.reset_states()
            avg_val_loss.reset_states()
            model.save_weights(
                'checkpoints/yolov3_train_{}.tf'.format(epoch))
    else:
        model.compile(optimizer=optimizer, loss=loss)

        callbacks = [
            ReduceLROnPlateau(verbose=1),
#            EarlyStopping(patience=3, verbose=1),
            ModelCheckpoint('checkpoints/yolov3_train_{epoch}.tf',
                            verbose=1, save_weights_only=True, save_best_only=True),
        ]

        model.fit(train_dataset,
                  epochs=FLAGS.epochs,
                  steps_per_epoch=steps_for_train,
                  callbacks=callbacks,
                  validation_data=val_dataset,
                  validation_steps=steps_for_val)

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass


from theano import tensor as T

from ..environment import BaseEnvironment
from ..objective import BaseObjective
from ..utils.format import check_list
from ..utils.logging import warn

class SessionBatchEnvironment(BaseEnvironment, BaseObjective):
    """
    A generic pseudo-environment that replays sessions defined on creation by theano expressions
    ignoring agent actions completely.

    The environment takes symbolic expression for sessions represented as (.observations, .actions, .rewards)
    Unlike SessionPoolEnvironment, this one does not store it's own pool of sessions.

    To create experience-replay sessions, call Agent.get_sessions with this as an environment.

    :param observations: a tensor or a list of tensors matching agent observation sequence [batch, tick, whatever]
    :type observations: theano tensor or a list of such
    :param single observation shapes: shapes of one-tick one-batch-item observations.
                            E.g. if lasagne shape is [None, 25(ticks), 3,210,160],
                            than single_observation_shapes must contain [3,210,160]
    :type single_observation_shapes: a list of tuples of integers
    :param actions: a tensor or a list of tensors matching agent actions sequence [batch, tick, whatever]
    :type actions: theano tensor or a list of such
    :param single action shapes: shapes of one-tick one-batch-item actions. Similar to observations.
                                    All scalar means that each action has shape (,),
                                    lasagne sequence layer being of shape (None, seq_length)
    :type single_observation_shapes: a list of tuples of integers
    :param rewards: a tensor matching agent rewards sequence [batch, tick]
    :type rewards: theano tensor

    :param is_alive: whether or not session has still not finished by a particular tick. Always alive by default.
    :type is_alive: theano tensor or None
    :param preceding_agent_memory: a tensor or a list of such storing what was in agent's memory prior
                                    to the first tick of the replay session.
    :type actions: theano tensor or a list of such

    How does it tick:

    During experience replay sessions,
     - observations, actions and rewards match original ones
     - agent memory states, Q-values and all in-agent expressions (but for actions) will correspond to what
       agent thinks NOW about the replay (not what it thought when he commited actions)
     - preceding_agent_memory [optional] - what was agent's memory state prior to the first tick of the replay session.

    Although it is possible to get rewards via the regular functions, it is usually faster to take self.rewards as rewards
    with no additional computation.

    """

    def __init__(self, observations, single_observation_shapes,
                 actions=None, single_action_shapes="all_scalar",
                 rewards=None, is_alive=None, preceding_agent_memories=None):

        # setting environmental variables. Their shape is [batch,time,something]
        self.observations = check_list(observations)
        self.single_observation_shapes = check_list(single_observation_shapes)

        for obs, obs_shape in zip(self.observations, self.single_observation_shapes):
            assert obs.ndim == len(obs_shape) + 2

        if actions is not None:
            self.actions = check_list(actions)
            self.single_action_shapes = single_action_shapes
            if self.single_action_shapes == "all_scalar":
                self.single_action_shapes = [tuple()] * len(self.actions)

        self.rewards = rewards
        self.is_alive = is_alive

        if preceding_agent_memories is not None:
            self.preceding_agent_memories = check_list(preceding_agent_memories)

        self.padded_observations = [
            T.concatenate([obs, T.zeros_like(obs[:, :1])], axis=1)
            for obs in self.observations
            ]

        self.batch_size = self.observations[0].shape[0]
        self.sequence_length = self.observations[0].shape[1]

        BaseEnvironment.__init__(self,
                                 state_shapes=tuple(()),
                                 observation_shapes=self.single_observation_shapes,
                                 action_shapes=self.single_action_shapes,
                                 state_dtypes=("int32",),
                                 observation_dtypes=tuple(obs.dtype for obs in self.observations),
                                 action_dtypes=tuple(act.dtype for act in self.actions)
                                 )

    def get_action_results(self, last_states, actions, **kwargs):
        """Computes environment state after processing agent's action.

        :param last_states: Environment state on previous tick.
        :type last_states: float[batch_id, memory_id0[, memory_id1,...]]

        :param actions: Agent action after observing last state.
        :type actions: int[batch_id]

        :returns:
            new_state float[batch_id, memory_id0[, memory_id1,...]]: Environment state after processing agent's action.
            observation float[batch_id, n_agent_inputs]: What agent observes after commiting the last action.
        """
        warn("Warning - a session pool has all the observations already stored as .observations property."
             "Recomputing them this way is probably just a slower way of calling your_session_pool.observations")

        time_i = check_list(last_states)[0]

        batch_range = T.arange(time_i.shape[0])

        new_observations = [obs[batch_range, time_i + 1] for obs in self.padded_observations]
        return [time_i + 1], new_observations

    def get_reward(self, session_states, session_actions, batch_id):
        """
        WARNING! this runs on a single session, not on a batch.
        Reward given for taking the action in current environment state.

        :param session_states: Environment state before taking action.
        :type session_states: float[batch_id, memory_id]

        :param session_actions: Agent action at this tick.
        :type session_actions: int[batch_id]

        :returns:
            reward float[batch_id]: Reward for taking action from the given state.
        """
        warn("Warning - a session pool has all the rewards already stored as .rewards property."
             "Recomputing them this way is probably just a slower way of calling your_session_pool.rewards")
        return self.rewards[batch_id, :]


######################################################################
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. #
# SPDX-License-Identifier: MIT-0                                     #
######################################################################

import numpy as np
import pandas as pd
from tqdm import tqdm
import torch
from torch.autograd import Variable
from ann_model import ANN_Model

def initialize(model_loc, gpu_mode, **kwargs):
    print('Initializing: %s' % (model_loc))
    cuda_available = torch.cuda.is_available()
    print('CUDA available? %s'%cuda_available)
    effective_gpu_mode = gpu_mode & cuda_available
    print('Effective GPU Mode: %s'%effective_gpu_mode)
    mlp_pt = ANN_Model(model_loc, gpu_mode)
    print(mlp_pt.named_modules)

    ret_dict = {
        'mlp_pt': mlp_pt,
        'gpu_mode': effective_gpu_mode
    }
    return ret_dict

# inp_dict has
# analyticSettings
# inputs
def run(mlp_pt, gpu_mode, inp_dict):
    analyticSettings = inp_dict['analyticSettings']
    inputs = inp_dict['inputs']
    bias_weight_array = None

    # Extract the weight_bias_map from modelParams
    bias_weight_map = analyticSettings['modelParams']['bias_weight_map']
    # if "bias_weight_array" in analyticSettings['modelParams']:
    #     bias_weight_array = np.array(analyticSettings['modelParams']['bias_weight_array']).astype(np.float32)

    # Override with any values present in the tunableParams
    if "tunableParams" in analyticSettings:
        if "bias_weight_array" in analyticSettings['tunableParams']:
            bias_weight_array = np.array(analyticSettings['tunableParams']['bias_weight_array']).astype(np.float32)
            for layer_name, weight_bias in bias_weight_map.items():
                # Extract the size of the weights and bias
                (loc_weight_array, loc_bias_array) = mlp_pt.get_layer_weights_bias(layer_name=layer_name)
                if 'weights' in weight_bias:
                    loc_array = bias_weight_array[np.array(weight_bias['weights'])]
                    # Reshape the array to the correct size of the weights array
                    mlp_pt.set_layer_weights(layer_name=layer_name, weights_np=loc_array.reshape(loc_weight_array.shape))
                if 'bias' in weight_bias:
                    loc_array = bias_weight_array[np.array(weight_bias['bias'])]
                    mlp_pt.set_layer_bias(layer_name=layer_name, bias_np=loc_array.reshape(1, -1))

    # Assemble the x_pred as numpy and convert it to a torch object running on a GPU
    # inputs = analyticSettings['inputs']
    # Use the modelParams input_var_order to assemble an array
    input_var_order = analyticSettings['modelParams']['input_var_order']
    n_cols = len(input_var_order)
    n_rows = len(inputs[input_var_order[0]])
    x_pred = np.zeros((n_rows, n_cols))
    for idx, cur_var in enumerate(input_var_order):
        x_pred[:, idx] = np.array(inputs[cur_var])
        # inputs_array_list.append(inputs[cur_var])
    # x_pred = np.hstack(inputs_array_list)
    # x_pred = np.array(inputs['x_pred'])
    print('Forward Pass: %s'%(str(x_pred.shape)))
    if (gpu_mode):
        x_pred_pt = Variable(torch.from_numpy(x_pred).type(torch.FloatTensor)).cuda()
    else:
        x_pred_pt = Variable(torch.from_numpy(x_pred).type(torch.FloatTensor))

    # Forward call detaches and converts to numpy
    out_pred_update = mlp_pt.forward(x_pred_pt)

    ret_dict = {}
    output_var_order = analyticSettings['modelParams']['output_var_order']
    for idx, cur_var in enumerate(output_var_order):
        ret_dict[cur_var] = out_pred_update[:, idx].tolist()

    return ret_dict

    #
    # ret_dict = {
    #     'y': out_pred_update.tolist()
    # }
    #
    # return ret_dict

def cleanup(**kwargs):
    return

if __name__ == "__main__":
    import os

    model_loc = os.getenv('MODEL_LOC')
    print('Initializing w/:%s' % (model_loc))
    gpu_mode = torch.cuda.is_available()
    init_dict = initialize(model_loc=model_loc, gpu_mode=gpu_mode)

    inp_dict = {}

    loc_analyticSettings = {}
    loc_analyticSettings['modelParams'] = {
        "input_var_order": ["x_1", "x_2", "x_3"],
        "output_var_order": ["y_1"],
        # "bias_weight_array": list(np.ones((5,))),
        "bias_weight_map": {
            "final": {
                "weights": [1, 2, 3, 4],
                "bias": [0]
            }
        }
    }
    loc_analyticSettings['tunableParams'] = {
        "bias_weight_array": list(np.random.random(5, ))
    }

    # Set the Inputs
    n_rows = 10
    inputs = {}
    for idx, cur_var in enumerate(loc_analyticSettings['modelParams']['input_var_order']):
        inputs[cur_var] = np.random.random(n_rows).tolist()
    #
    # inputs = {
    #     "x_pred": list(np.random.random((10, 3)))
    # }

    inp_dict['analyticSettings'] = loc_analyticSettings
    inp_dict['inputs'] = inputs

    print('Simulating')
    ret_dict = run(inp_dict=inp_dict,
                   **init_dict)

    print(ret_dict)

    print('Clean Up')
    cleanup(**init_dict)


"""
Wrapper classes for internal ML models.

MultichannelPipelines treat all internal component as transfomers (i.e.
invoking fit/transform/fit_transform).  As a consequence, when predictors are
used internally (e.g. for voting or stacking) a transformer interface must be
added to the internal predictors.  In practice, this means choosing a
prediction method to use when transforming, converting 1D outputs to 2D
outputs, and applying internal cross validation training when required.

:class:`SingleChannel` and :class:`Multichannel` classes add a transformer
interface to single channel and multichannel predictors respectively.

:class:`SingleChannelCV` and :class:`MultichannelCV` classes add a transformer
interface and internal cross validaiton training to single channel and
multichannel predictors respectively.  Internal cross validation (internal cv)
training is typically used when outputs of a base predictor will be used to
train a meta-predictor.  It guarantees that base predictors do not make
inferences on their own training samples (1).  Internal cv training can improve
meta-predictor accuracy if overfitting is a limiting problem, or it can reduce
metapredictor accuracy if the number of training samples is limiting.

(1) Wolpert, David H. "Stacked generalization." Neural networks 5.2
    (1992): 241-259.
"""

import functools
import numpy as np
from sklearn.metrics import log_loss

import pipecaster.utils as utils
import pipecaster.config as config
from pipecaster.utils import Cloneable, Saveable
from pipecaster.cross_validation import cross_val_predict, score_predictions

__all__ = ['make_transformer', 'make_cv_transformer', 'unwrap_predictor',
           'unwrap_model']

def make_transformer(predictor, transform_method='auto'):
    """
    Add transform methods to a predictor.

    Parameters
    ----------
    predictor : scikit-learn predictor or multichannel predictor
        Predictor to wrap.
    transform_method : str, default='auto'
        - Name of the prediction method to call when transforming (e.g. when
          outputting meta-features).
        - If 'auto' :
            - If classifier : method picked using
              config.transform_method_precedence order (default:
              predict_proba->predict_log_proba->decision_function->predict).
            - If regressor : 'predict'

    Returns
    -------
    Predictor/transformer
        A wrapped predictor with both predictor and transformer interfaces.

    Examples
    --------
    ::

        from sklearn.ensemble import GradientBoostingClassifier
        import pipecaster as pc

        Xs, y, X_types = pc.make_multi_input_classification(n_informative_Xs=3,
                                                            n_random_Xs=2)
        clf = pc.MultichannelPipeline(n_channels=5)
        base_clf = GradientBoostingRegressor()
        base_clf = pc.make_transformer(base_clf)
        clf.add_layer(base_clf)
        clf.add_layer(pc.SoftVotingClassifier())
        pc.cross_val_score(clf, Xs, y, cv=3)
        # output: [0.8529411764705882, 0.9411764705882353, 0.96875]
    """
    if utils.is_multichannel(predictor):
        return Multichannel(predictor, transform_method)
    else:
        return SingleChannel(predictor, transform_method)

def make_cv_transformer(predictor, transform_method='auto', internal_cv=5,
             score_method='auto', scorer='auto', cv_processes=1):
    """
    Add internal cross validation training and transform methods to a
    predictor.

    Parameters
    ----------
    predictor : scikit-learn predictor or multichannel predictor
        Predictor to wrap.
    transform_method : str, default='auto'
        - Name of the prediction method to call when transforming (e.g. when
          outputting meta-features).
        - If 'auto' :
            - If classifier : method picked using
              config.transform_method_precedence order (default:
              predict_proba->predict_log_proba->decision_function->predict).
            - If regressor : 'predict'
    internal_cv : int, None, or callable, default=5
        - Function for train/test subdivision of the training data.  Used to
          estimate performance of base classifiers and ensure they do not
          generate predictions from their training samples during
          meta-predictor training.
        - If int > 1: StratifiedKfold(n_splits=internal_cv) if classifier or
          KFold(n_splits=internal_cv) if regressor.
        - If {None, 1}: disable internal cv.
        - If callable: Assumed to be split generator like scikit-learn KFold.
    score_method : str, default='auto'
        - Name of prediction method used when scoring predictor performance.
        - If 'auto' :
            - If classifier : method picked using
              config.score_method_precedence order (default:
              ppredict_proba->predict_log_proba->decision_function->predict).
            - If regressor : 'predict'
    scorer : callable, default='auto'
        Callable that computes a figure of merit score for the internal_cv run.
        The score is exposed as score_ attribute during fit_transform().
        - If 'auto':
            - explained_variance_score for regressors with predict()
            - roc_auc_score for classifiers with {predict_proba,
              predict_log_proba, decision_function}
            - balanced_accuracy_score for classifiers with only predict()
        - If callable: A scorer with signature: score = scorer(y_true, y_pred).
    cv_processes : int or 'max', default=1
        - The number of parallel processes to run for internal cross
          validation.
        - If int : Use up to cv_processes number of processes.
        - If 'max' : Use all available CPUs.

    Returns
    -------
    Predictor/transformer
        A wrapped predictor with both predictor and transformer interfaces.
        Internal cross_validation training occurs during calls to
        fit_transform().

    Examples
    --------
    ::

        from sklearn.ensemble import GradientBoostingClassifier
        import pipecaster as pc

        Xs, y, X_types = pc.make_multi_input_classification(n_informative_Xs=3,
                                                      n_random_Xs=2)
        clf = pc.MultichannelPipeline(n_channels=5)
        base_clf = GradientBoostingRegressor()
        base_clf = pc.make_cv_transformer(base_clf)
        clf.add_layer(base_clf)
        clf.add_layer(pc.MultichannelPredictor(GradientBoostingClassifier()))
        pc.cross_val_score(clf, Xs, y, cv=3)
        # output: [0.8529411764705882, 0.9080882352941176, 1.0]
    """
    if utils.is_multichannel(predictor):
        return MultichannelCV(predictor, transform_method, internal_cv,
                              score_method, scorer, cv_processes)
    else:
        return SingleChannelCV(predictor, transform_method, internal_cv,
                              score_method, scorer, cv_processes)

class SingleChannel(Cloneable, Saveable):
    """
    Add transformer interface to a scikit-learn predictor.

    Wrapper class that provides scikit-learn conformant predictors with
    transform() and fit_transform methods().

    Parameters
    ----------
    predictor : predictor instance
        The scikit-learn conformant estimator/predictor to wrap.
    transform_method : str, default='auto'
        - Name of the prediction method to call when transforming (e.g. when
          outputting meta-features).
        - If 'auto' :
            - If classifier : method picked using
              config.transform_method_precedence order (default:
              predict_proba->predict_log_proba->decision_function->predict).
            - If regressor : 'predict'

    Examples
    --------
    Model stacking, classification:
    ::

        from sklearn.ensemble import GradientBoostingClassifier
        from sklearn.svm import SVC
        import pipecaster as pc

        Xs, y, _ = pc.make_multi_input_classification(n_informative_Xs=3,
                                                      n_random_Xs=7)
        clf = pc.MultichannelPipeline(n_channels=10)
        base_clf = GradientBoostingClassifier()
        base_clf = pc.transform_wrappers.SingleChannel(base_clf)
        clf.add_layer(base_clf, pipe_processes='max')
        clf.add_layer(pc.MultichannelPredictor(SVC()))
        pc.cross_val_score(clf, Xs, y, cv=3)
        # output: [0.8529411764705882, 0.8216911764705883, 0.9099264705882353]

    Model stacking, regression:
    ::
        from sklearn.ensemble import GradientBoostingRegressor
        from sklearn.svm import SVR
        import pipecaster as pc

        Xs, y, _ = pc.make_multi_input_regression(n_informative_Xs=7,
                                                  n_random_Xs=3)
        clf = pc.MultichannelPipeline(n_channels=10)
        base_clf = GradientBoostingRegressor()
        base_clf = pc.transform_wrappers.SingleChannel(base_clf)
        clf.add_layer(base_clf, pipe_processes=1)
        clf.add_layer(pc.MultichannelPredictor(SVR()))
        pc.cross_val_score(clf, Xs, y, cv=3)
        # output: [0.077183453, 0.067682880449, 0.07849665]

    Notes
    -----
    This class uses reflection to expose the predictor methods found in the
    object that it wraps, so the method attributes in a SingleChannel instance
    are not usually identical to the method attributes of the SingleChannel
    class.

    If a sample_weight parameter is sent to the fit() method but the wrapped
    predictor doesn't accept this argument, fit() will be called without the
    sample_weight parameter and no warning will be given.
    """

    def __init__(self, predictor, transform_method='auto'):
        self._params_to_attributes(SingleChannel.__init__, locals())
        utils.enforce_fit(predictor)
        utils.enforce_predict(predictor)
        self._add_predictor_interface(predictor)
        self._set_estimator_type(predictor)

    def _set_estimator_type(self, predictor):
        if hasattr(predictor, '_estimator_type') is True:
            self._estimator_type = predictor._estimator_type

    def _add_predictor_interface(self, predictor):
        for method_name in config.recognized_pred_methods:
            if hasattr(predictor, method_name):
                prediction_method = functools.partial(self.predict_with_method,
                                                      method_name=method_name)
                setattr(self, method_name, prediction_method)

    def _add_model_interface(self, model, X):
        detected_methods = utils.detect_predict_methods(model, X)
        for method_name in detected_methods:
            prediction_method = functools.partial(self.predict_with_method,
                                                  method_name=method_name)
            setattr(self, method_name, prediction_method)

    def _remove_predictor_interface(self):
        for method_name in config.recognized_pred_methods:
            if hasattr(self, method_name):
                delattr(self, method_name)

    def set_transform_method(self, method_name):
        self.transform_method = method_name
        return self

    def get_transform_method(self):
        if self.transform_method == 'auto':
            method_name = utils.get_transform_method(self)
            if method_name is None:
                raise NameError('model lacks a recognized method for '
                                'conversion to transformer')
        else:
            method_name = self.transform_method

        return method_name

    def fit(self, X, y=None, **fit_params):
        self.model = utils.get_clone(self.predictor)
        is_classifier = utils.is_classifier(self.predictor)
        if y is None:
            try:
                self.model.fit(X, **fit_params)
            except:
                self.model.fit(X)
        else:
            if is_classifier:
                self.classes_, y = np.unique(y, return_inverse=True)
            try:
                self.model.fit(X, y, **fit_params)
            except:
                self.model.fit(X, y)
        self._set_estimator_type(self.model)
        self._remove_predictor_interface()
        self._add_model_interface(self.model, X)

        return self

    def predict_with_method(self, X, method_name):
        if hasattr(self, 'model') is False:
            raise utils.FitError('prediction attempted before model fitting')
        if hasattr(self.model, method_name):
            predict_method = getattr(self.model, method_name)
            predictions = predict_method(X)
        else:
            raise NameError('prediction method {} not found in {} attributes'
                            .format(method_name, self.model))
        if utils.is_classifier(self) and method_name == 'predict':
            predictions = self.classes_[predictions]

        return predictions

    def transform(self, X):
        if hasattr(self, 'model'):
            transformer = getattr(self.model, self.get_transform_method())
            X_t = transformer(X)
            # convert output array to output matrix:
            if len(X_t.shape) == 1:
                X_t = X_t.reshape(-1, 1)
            # drop redundant prob output from binary classifiers:
            elif (len(X_t.shape) == 2 and X_t.shape[1] == 2 and
                  utils.is_classifier(self.model)):
                X_t = X_t[:, 1].reshape(-1, 1)
            return X_t
        else:
            raise utils.FitError('transform called before model fitting')

    def fit_transform(self, X, y=None, **fit_params):
        self.fit(X, y, **fit_params)
        return self.transform(X)

    def _more_tags(self):
        return {'multichannel': False}

    def get_clone(self):
        """
        Get a stateful clone.
        """
        clone = super().get_clone()
        if hasattr(self, 'classes_'):
            clone.classes_ = self.classes_.copy()
        if hasattr(self, 'model'):
            clone.model = utils.get_clone(self.model)
            clone._set_estimator_type(self.model)
            clone._remove_predictor_interface()
            clone._add_predictor_interface(self)
        return clone

    def get_descriptor(self, verbose=1):
        return '{' + utils.get_descriptor(self.predictor, verbose,
                                          self.get_params()) + '}tr'

class SingleChannelCV(SingleChannel):
    """
    Add transformer interface and internal cross validation training to
    scikit-learn predictor.

    Wrapper class that provides predictors with transform() and fit_transform()
    methods, and internal cross validation training with performance scoring.

    Parameters
    ----------
    predictor : predictor instance
        The scikit-learn conformant predictor to wrap.
    transform_method : str, default='auto'
        - Name of the prediction method to call when transforming (e.g. when
          outputting meta-features).
        - If 'auto' :
            - If classifier : method picked using
              config.transform_method_precedence order (default:
              predict_proba->predict_log_proba->decision_function->predict).
            - If regressor : 'predict'
    internal_cv : int, None, or callable, default=5
        - Function for train/test subdivision of the training data.  Used to
          estimate performance of base classifiers and ensure they do not
          generate predictions from their training samples during
          meta-predictor training.
        - If int > 1: StratifiedKfold(n_splits=internal_cv) if classifier or
          KFold(n_splits=internal_cv) if regressor.
        - If {None, 1}: disable internal cv.
        - If callable: Assumed to be split generator like scikit-learn KFold.
    score_method : str, default='auto'
        - Name of prediction method used when scoring predictor performance.
        - If 'auto' :
            - If classifier : method picked using
              config.score_method_precedence order (default:
              ppredict_proba->predict_log_proba->decision_function->predict).
            - If regressor : 'predict'
    scorer : callable, default='auto'
        Callable that computes a figure of merit score for the internal_cv run.
        The score is exposed as score_ attribute during fit_transform().
        - If 'auto':
            - explained_variance_score for regressors with predict()
            - roc_auc_score for classifiers with {predict_proba,
              predict_log_proba, decision_function}
            - balanced_accuracy_score for classifiers with only predict()
        - If callable: A scorer with signature: score = scorer(y_true, y_pred).
    cv_processes : int or 'max', default=1
        - The number of parallel processes to run for internal cross
          validation.
        - If int : Use up to cv_processes number of processes.
        - If 'max' : Use all available CPUs.

    Examples
    --------
    Model stacking:
    ::

        from sklearn.ensemble import GradientBoostingClassifier
        from sklearn.svm import SVC
        import pipecaster as pc

        Xs, y, _ = pc.make_multi_input_classification(n_informative_Xs=3,
                                                      n_random_Xs=7)
        clf = pc.MultichannelPipeline(n_channels=10)
        base_clf = GradientBoostingClassifier()
        base_clf = pc.transform_wrappers.SingleChannelCV(base_clf)
        clf.add_layer(base_clf, pipe_processes='max')
        clf.add_layer(pc.MultichannelPredictor(SVC()))
        pc.cross_val_score(clf, Xs, y, cv=3)
        # output: [0.9411764705882353, 0.8897058823529411, 0.9963235294117647]

    Notes
    -----
    fit().transform() is not the same as fit_tranform() because only the latter
    uses internal cv training and inference.
    On calls to fit_transform() the model is fit on both the entire training
    set and cv splits of the training set. The model fit on the entire dataset
    is stored for futer inference.  The models fit on cv splits are used
    to make the outputs of fit_transform() but are not stored for future use.

    This class uses reflection to expose the predictor methods found in the
    object that it wraps, so the method attributes in a SingleChannelCV
    instance are usually not identical to the method attributes of the
    SingleChannelCV class.

    If a sample_weight parameter is sent to the fit() method but the wrapped
    predictor doesn't accept this argument, fit() will be called without the
    sample_weight parameter and no warning will be given.
    """

    def __init__(self, predictor, transform_method='auto', internal_cv=5,
                 score_method='auto', scorer='auto', cv_processes=1):
        self._params_to_attributes(SingleChannelCV.__init__, locals())
        super().__init__(predictor, transform_method)

    def fit_transform(self, X, y=None, groups=None, **fit_params):
        is_classifier = utils.is_classifier(self.predictor)
        if y is not None and is_classifier:
            self.classes_, y = np.unique(y, return_inverse=True)

        self.fit(X, y, **fit_params)
        transform_method = self.get_transform_method()

        # if internal cv training is disabled
        if (self.internal_cv is None or
                (type(self.internal_cv) == int and self.internal_cv < 2)):
            y_transform = self.transform(X)
        # internal cv training is enabled
        else:
            split_results = cross_val_predict(self.predictor, X, y,
                                    groups=groups,
                                    predict_method=None,
                                    transform_method=transform_method,
                                    score_method=self.score_method,
                                    cv=self.internal_cv, combine_splits=True,
                                    n_processes=self.cv_processes,
                                    fit_params=fit_params)

            y_transform = split_results['transform']['y_pred']
            y_score = split_results['score']['y_pred']
            is_binary = (True if is_classifier and len(self.classes_) == 2
                         else False)

            score_method = split_results['score']['method']
            self.score_ = score_predictions(y, y_score, score_method,
                                            self.scorer, is_classifier,
                                            is_binary)

        # convert output array to output matrix:
        X_t = y_transform
        if len(X_t.shape) == 1:
            X_t = X_t.reshape(-1, 1)

        # drop the redundant prob output from binary classifiers:
        elif (len(X_t.shape) == 2 and X_t.shape[1] == 2 and
              utils.is_classifier(self.model)):
            X_t = X_t[:, 1].reshape(-1, 1)

        return X_t

    def _more_tags(self):
        return {'multichannel': False}

    def get_descriptor(self, verbose=1):
        return '{' + utils.get_descriptor(self.predictor, verbose,
                                          self.get_params()) + '}cvtr'

    def get_clone(self):
        clone = super().get_clone()
        if hasattr(self, 'scores_'):
            clone.scores_ = self.scores_
        return clone

class Multichannel(Cloneable, Saveable):
    """
    Add transformer interface to a multichannel predictor.

    Wrapper class that provides pipecaster's multichannel predictors with
    transform() and fit_transform methods().

    Parameters
    ----------
    multichannel_predictor : multichannel predictor instance
        The predictor to wrap.
    transform_method : str, default='auto'
        - Name of the prediction method to call when transforming (e.g. when
          outputting meta-features).
        - If 'auto' :
            - If classifier : method picked using
              config.transform_method_precedence order (default:
              predict_proba->predict_log_proba->decision_function->predict).
            - If regressor : 'predict'

    Examples
    --------
    model stacking:
    ::

        from sklearn.ensemble import GradientBoostingClassifier
        from sklearn.svm import SVC
        import pipecaster as pc

        Xs, y, _ = pc.make_multi_input_classification(n_informative_Xs=3,
                                                      n_random_Xs=7)
        clf = pc.MultichannelPipeline(n_channels=10)
        base_clf = pc.MultichannelPredictor(GradientBoostingClassifier())
        base_clf = pc.transform_wrappers.Multichannel(base_clf)
        clf.add_layer(5, base_clf, 5, base_clf, pipe_processes=1)
        clf.add_layer(pc.MultichannelPredictor(SVC()))
        pc.cross_val_score(clf, Xs, y, cv=3)
        # output: [0.9411764705882353, 0.9411764705882353, 0.8768382352941176]

    Notes
    -----
    This class uses reflection to expose the predictor methods found in the
    object that it wraps, so the method attributes in a Multichannel instance
    are not usually identical to the method attributes of the Multichannel
    class.
    """

    def __init__(self, multichannel_predictor, transform_method='auto'):
        self._params_to_attributes(Multichannel.__init__, locals())
        utils.enforce_fit(multichannel_predictor)
        utils.enforce_predict(multichannel_predictor)
        utils.enforce_predict(multichannel_predictor)
        self._add_predictor_interface(multichannel_predictor)
        self._set_estimator_type(multichannel_predictor)

    def _set_estimator_type(self, predictor):
        if hasattr(predictor, '_estimator_type') is True:
            self._estimator_type = predictor._estimator_type

    def _add_predictor_interface(self, predictor):
        for method_name in config.recognized_pred_methods:
            if hasattr(predictor, method_name):
                prediction_method = functools.partial(self.predict_with_method,
                                                      method_name=method_name)
                setattr(self, method_name, prediction_method)

    def _add_model_interface(self, model, Xs):
        detected_methods = utils.detect_predict_methods(model, Xs)
        for method_name in detected_methods:
            prediction_method = functools.partial(self.predict_with_method,
                                                  method_name=method_name)
            setattr(self, method_name, prediction_method)

    def _remove_predictor_interface(self):
        for method_name in config.recognized_pred_methods:
            if hasattr(self, method_name):
                delattr(self, method_name)

    def get_transform_method(self):
        if self.transform_method == 'auto':
            if hasattr(self, 'model'):
                method_name = utils.get_transform_method(self.model)
            else:
                method_name = utils.get_transform_method(
                                                self.multichannel_predictor)
            if method_name is None:
                raise NameError('model lacks a recognized method for '
                                'conversion to transformer')
        else:
            method_name = self.transform_method

        return method_name

    def fit(self, Xs, y=None, **fit_params):
        self.model = utils.get_clone(self.multichannel_predictor)
        if y is None:
            self.model.fit(Xs, **fit_params)
        else:
            if utils.is_classifier(self.model):
                self.classes_, y = np.unique(y, return_inverse=True)
            self.model.fit(Xs, y, **fit_params)

        self._set_estimator_type(self.model)
        self._remove_predictor_interface()
        self._add_model_interface(self.model, Xs)

        return self

    def predict_with_method(self, Xs, method_name):
        if hasattr(self, 'model') is False:
            raise FitError('prediction attempted before call to fit()')
        prediction_method = getattr(self.model, method_name)
        predictions = prediction_method(Xs)
        if utils.is_classifier(self) and method_name == 'predict':
            predictions = self.classes_[predictions]
        return predictions

    def transform(self, Xs):
        if hasattr(self, 'model') is False:
            raise FitError('transform attempted before call to fit()')
        tansformer = getattr(self.model, self.get_transform_method())
        predictions = np.array(tansformer(Xs))

        # convert output array to output matrix:
        if len(predictions.shape) == 1:
            predictions = predictions.reshape(-1, 1)
        # drop the redundant prob output from binary classifiers:
        elif (len(predictions.shape) == 2 and predictions.shape[1] == 2
              and utils.is_classifier(self.model)):
            predictions = predictions[:, 1].reshape(-1, 1)

        Xs_t = [predictions if i == 0 else None for i, X in enumerate(Xs)]
        return Xs_t

    def fit_transform(self, Xs, y=None, **fit_params):
        self.fit(Xs, y, **fit_params)
        return self.transform(Xs)

    def get_descriptor(self, verbose=1):
        return '{' + utils.get_descriptor(self.multichannel_predictor, verbose,
                                          self.get_params()) + '}tr'

    def get_clone(self):
        """
        Get a stateful clone.
        """
        clone = super().get_clone()
        if hasattr(self, 'classes_'):
            clone.classes_ = self.classes_.copy()
        if hasattr(self, 'model'):
            clone.model = utils.get_clone(self.model)
            clone._set_estimator_type(self.model)
            clone._remove_predictor_interface()
            clone._add_predictor_interface(self)
        return clone

class MultichannelCV(Multichannel):
    """
    Add transformer interface and internal cross validation training to
    multichannel predictor.

    Wrapper class that provides pipecaster's multichannel predictors with
    transform() and fit_transform() methods, and internal cross validation
    training with performance scoring.

    Parameters
    ----------
    multichannel_predictor : multichannel_predictor instance
        The pipecaster predictor to wrap.
    transform_method : str, default='auto'
        - Name of the prediction method to call when transforming (e.g. when
          outputting meta-features).
        - If 'auto' :
            - If classifier : method picked using
              config.transform_method_precedence order (default:
              predict_proba->predict_log_proba->decision_function->predict).
            - If regressor : 'predict'
    internal_cv : int, None, or callable, default=5
        - Function for train/test subdivision of the training data.  Used to
          estimate performance of base classifiers and ensure they do not
          generate predictions from their training samples during
          meta-predictor training.
        - If int > 1: StratifiedKfold(n_splits=internal_cv) if classifier or
          KFold(n_splits=internal_cv) if regressor.
        - If {None, 1}: disable internal cv.
        - If callable: Assumed to be split generator like scikit-learn KFold.
    score_method : str, default='auto'
        - Name of prediction method used when scoring predictor performance.
        - If 'auto' :
            - If classifier : method picked using
              config.score_method_precedence order (default:
              ppredict_proba->predict_log_proba->decision_function->predict).
            - If regressor : 'predict'
    scorer : callable, default='auto'
        Callable that computes a figure of merit score for the internal_cv run.
        The score is exposed as score_ attribute during fit_transform().
        - If 'auto':
            - explained_variance_score for regressors with predict()
            - roc_auc_score for classifiers with {predict_proba,
              predict_log_proba, decision_function}
            - balanced_accuracy_score for classifiers with only predict()
        - If callable: A scorer with signature: score = scorer(y_true, y_pred).
    cv_processes : int or 'max', default=1
        - The number of parallel processes to run for internal cross
          validation.
        - If int : Use up to cv_processes number of processes.
        - If 'max' : Use all available CPUs.

    Examples
    --------
    model stacking:
    ::

        from sklearn.ensemble import GradientBoostingClassifier
        from sklearn.svm import SVC
        import pipecaster as pc

        Xs, y, _ = pc.make_multi_input_classification(n_informative_Xs=3,
                                                      n_random_Xs=7)
        clf = pc.MultichannelPipeline(n_channels=10)
        base_clf = pc.MultichannelPredictor(GradientBoostingClassifier())
        base_clf = pc.transform_wrappers.MultichannelCV(base_clf)
        clf.add_layer(5, base_clf, 5, base_clf, pipe_processes=1)
        clf.add_layer(pc.MultichannelPredictor(SVC()))
        pc.cross_val_score(clf, Xs, y, cv=3)
        # output: [0.8823529411764706, 0.9393382352941176, 0.9080882352941176]

    Notes
    -----
    fit().transform() is not the same as fit_tranform() because only the latter
    uses internal cv training and inference.
    On calls to fit_transform() the model is fit on both the entire training
    set and cv splits of the training set. The model fit on the entire dataset
    is stored for future inferences.  The models fit on cv splits are used
    to make the outputs of fit_transform() but are not stored for future use.

    This class uses reflection to expose the predictor methods found in the
    object that it wraps, so the method attributes in a MultichannelCV
    instance are usually not identical to the method attributes of the
    MultichannelCV class.

    If a sample_weight parameter is sent to the fit() method but the wrapped
    predictor doesn't accept this argument, fit() will be called without the
    sample_weight parameter and no warning will be given.
    """

    def __init__(self, multichannel_predictor, transform_method='auto',
                 internal_cv=5, score_method='auto',
                 scorer='auto', cv_processes=1):
        self._params_to_attributes(MultichannelCV.__init__, locals())
        super().__init__(multichannel_predictor, transform_method)

    def fit_transform(self, Xs, y=None, groups=None, **fit_params):
        is_classifier = utils.is_classifier(self.multichannel_predictor)
        if y is not None and is_classifier:
            self.classes_, y = np.unique(y, return_inverse=True)

        self.fit(Xs, y, **fit_params)
        transform_method = self.get_transform_method()

        # if internal cv training is disabled
        if (self.internal_cv is None or
                (type(self.internal_cv) == int and self.internal_cv < 2)):
            y_transform = self.transform(X)
        # internal cv training is enabled
        else:
            split_results = cross_val_predict(self.multichannel_predictor,
                                    Xs, y,
                                    groups=groups,
                                    predict_method=None,
                                    transform_method=transform_method,
                                    score_method=self.score_method,
                                    cv=self.internal_cv, combine_splits=True,
                                    n_processes=self.cv_processes,
                                    fit_params=fit_params)

            y_transform = split_results['transform']['y_pred']
            y_score = split_results['score']['y_pred']
            is_binary = (True if is_classifier and len(self.classes_) == 2
                         else False)
            score_method = split_results['score']['method']
            self.score_ = score_predictions(y, y_score, score_method,
                                            self.scorer, is_classifier,
                                            is_binary)

        # convert predictions to transformed matrix:
        X_t = y_transform
        if len(X_t.shape) == 1:
            X_t = X_t.reshape(-1, 1)

        # drop the redundant prob output from binary classifiers:
        elif (len(X_t.shape) == 2 and X_t.shape[1] == 2 and
              utils.is_classifier(self.model)):
            X_t = X_t[:, 1].reshape(-1, 1)

        Xs_t = [None for X in Xs]
        Xs_t[0] = X_t

        return Xs_t

    def get_descriptor(self, verbose=1):
        return '{' + utils.get_descriptor(self.multichannel_predictor, verbose,
                                          self.get_params()) + '}cvtr'

    def get_clone(self):
        """
        Get a stateful clone.
        """
        clone = super().get_clone()
        if hasattr(self, 'score_'):
            clone.score_ = self.score_
        return clone

def unwrap_predictor(pipe):
    """
    Return a predictor that is wrapped in a transform wrapper.
    """
    if type(pipe) not in [SingleChannel, SingleChannelCV, Multichannel,
                          MultichannelCV]:
        return pipe
    if type(pipe) in [Multichannel, MultichannelCV]:
        return pipe.multichannel_predictor
    else:
        return pipe.predictor

def unwrap_model(pipe):
    """
    Return a model that is wrapped in a transform wrapper.
    """
    if type(pipe) not in [SingleChannel, SingleChannelCV, Multichannel,
                          MultichannelCV]:
        return pipe
    if hasattr(pipe, 'model') is True:
        return pipe.model
    else:
        raise utils.FitError('no model found')


import numpy as np
import cv2
import tensorflow as tf

face_detection = cv2.CascadeClassifier('haar_cascade_face_detection.xml')

camera = cv2.VideoCapture(0)
camera.set(cv2.CAP_PROP_FRAME_WIDTH, 1024)
camera.set(cv2.CAP_PROP_FRAME_HEIGHT, 768)
settings = {
    'scaleFactor': 1.3,
    'minNeighbors': 5,
    'minSize': (50, 50)
}

labels = ["Neutral","Happy","Sad","Surprise","Angry"]

model = tf.keras.models.load_model('expression.model')

while True:
	ret, img = camera.read()
	gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
	detected = face_detection.detectMultiScale(gray, **settings)

	for x, y, w, h in detected:
		cv2.rectangle(img, (x, y), (x+w, y+h), (245, 135, 66), 2)
		cv2.rectangle(img, (x, y), (x+w//3, y+20), (245, 135, 66), -1)
		face = gray[y+5:y+h-5, x+20:x+w-20]
		face = cv2.resize(face, (48,48))
		face = face/255.0

		predictions = model.predict(np.array([face.reshape((48,48,1))])).argmax()
		state = labels[predictions]
		font = cv2.FONT_HERSHEY_SIMPLEX
		cv2.putText(img,state,(x+10,y+15), font, 0.5, (255,255,255), 2, cv2.LINE_AA)

	cv2.imshow('Facial Expression', img)
	img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
	# writer.writeFrame(img)
	if cv2.waitKey(5) != -1:
		break

camera.release()
cv2.destroyAllWindows()


# This Python file uses the following encoding: utf-8
from app.package.models.ActualProjectModel import ActualProjectModel
import os
import numpy as np
import sounddevice as sd

from PySide2.QtWidgets import QMainWindow
from PySide2.QtCore import QObject, Signal, QThread, Slot

from ..controllers.Calibrate_controller import CalibrateController
from ..models.Calibrate_model import CalibrateModel
from ..ui.Calibrate_ui import Ui_MainWindow as Calibrate_ui
from ..services.file import save_np_to_txt

def get_documents_dir():
    import platform
    if platform.system() == 'Windows':
        import ctypes.wintypes
        CSIDL_PERSONAL = 5       # My Documents
        SHGFP_TYPE_CURRENT = 0   # Get current, not default value

        buf = ctypes.create_unicode_buffer(ctypes.wintypes.MAX_PATH)
        ctypes.windll.shell32.SHGetFolderPathW(
            None, CSIDL_PERSONAL, None, SHGFP_TYPE_CURRENT, buf)

        return buf.value

class CalibrationWorker(QObject):
    finished = Signal()
    update_status = Signal(str)

    @staticmethod
    def log(msg: str) -> None:
        print(f'[CalibrationWorker] {msg}')

    def process(self, targetSpl):
        self.log('Processing')
        self.log(targetSpl)

        t = 5
        sd.default.samplerate = fs = 44100
        self.log(f'Time rec: {t}')

        # x = sd.rec(int(t * fs), channels=2)

        x = sd.rec(frames=int(t*fs), channels=1)
        sd.wait()
        x = np.sum(x, 1)  # change

        # print(np.max(x))
        # rms = np.sqrt(np.mean(np.square(x)))
        spl = 20*np.log10(np.std(x)/2e-5)

        self.save_calibration_file(spl, targetSpl)
        self.finished.emit()

    def save_calibration_file(self, actual_spl, expected_spl):
        self.log(f'Got: {actual_spl} dB')
        self.log(f'Expected: {expected_spl} dB')

        ActualProjectModel.calibration = {
            'actual': actual_spl, 'expected': expected_spl
        }

        path = os.path.join(get_documents_dir(), 'Scan&Paint Clone')
        file_name = 'calibration.dat'
        save_np_to_txt(np.array([expected_spl, actual_spl]), path, file_name)

class CalibrateView(QMainWindow, Calibrate_ui):

    @staticmethod
    def log(msg):
        print(f'[Calibrate/View] {msg}')

    def __init__(self, model: CalibrateModel, controller: CalibrateController):
        super(CalibrateView, self).__init__()
        self._model = model
        self._controller = controller

        self.setupUi(self)
        self.connect_to_controller()
        self.connect_to_model()
        self.set_default_values()

    def open(self):
        self.show()

    def close(self):
        self.hide()

    def connect_to_controller(self):
        self.calibrate_button.clicked.connect(self.start_calibration)
        self.go_back_button.clicked.connect(self._controller.go_back)
        pass

    def connect_to_model(self):
        pass

    def set_default_values(self):
        title = "Quick Scan: Calibrate"
        self.setWindowTitle(title)

    # region SLOTS

    @Slot()
    def start_calibration(self):
        try:
            spl = float(self.spl_value.text())
        except Exception:
            self.log(f'Could not convert {self.spl_value.text()} to int')

        # Create worker and thread
        worker = CalibrationWorker()
        thread = QThread()
        worker.moveToThread(thread)

        # Conect the 'started' signal of the Thread
        # to the corresponding function in the worker
        thread.started.connect(lambda: worker.process(spl))
        worker.update_status.connect(self.handle_update_status)

        # Connect the finished signal of thread & worker
        # with the corresponding stuff
        worker.finished.connect(thread.quit)
        worker.finished.connect(self.go_back)
        worker.finished.connect(worker.deleteLater)
        thread.finished.connect(thread.deleteLater)
        self.thread = thread
        self.thread.start()

    @Slot(str)
    def handle_update_status(self, value):
        print(value)

    @Slot()
    def go_back(self):
        self._controller._navigator.navigate('new_project')

    # endregion


# Copyright (c) 2020 Shubh Pachchigar
"""
EEG Data is taken from DEAP
The training data was taken from DEAP.
See my:
- Github profile: https://github.com/shubhe25p
- Email: shubhpachchigar@gmail.com
"""

import csv
import numpy as np
import scipy.spatial as ss
import scipy.stats as sst
import _pickle as cPickle
from pathlib import Path
import itertools
import random
import cv2

sampling_rate = 128
number_of_channel = 32 #considering only head electrodes
eeg_in_second = 63 #length of each trial
number_of_eeg = sampling_rate*eeg_in_second #total inputs from a single channel

channel_names=['Fp1','AF3','F3','F7','FC5','FC1','C3','T7','CP5','CP1','P3','P7','PO3','O1','Oz','Pz','Fp2','AF4','Fz','F4','F8','FC6','FC2','Cz','C4','T8','CP6','CP2','P4','P8','PO4','O2']

class predictEmotion(object):
	"""
	Receives EEG data preprocessing and predict emotion.
	"""

	# path is set to training data directory
	def __init__(self):
		"""
		Initializes training data and their classes.
		"""
		self.train_arousal = self.get_csv("train_std.csv")
		self.train_valence = self.get_csv("train_std.csv")
		self.class_arousal = self.get_csv("class_arousal.csv")
		self.class_valence = self.get_csv("class_valence.csv")

	def get_csv(self,path):
		"""
		Get data from csv and convert them to numpy python.
		Input: Path csv file.
		Output: Numpy array from csv data.
		"""
		#Get csv data to list
		file_csv = open(path)
		data_csv = csv.reader(file_csv)
		#convert list to array with a specific dtype
		data_training = np.array(list(data_csv),dtype=np.float64)
		return data_training

	def do_fft(self,all_channel_data):
		"""
		Do fft in each channel for all channels.
		Input: Channel data with dimension N x M. N denotes number of channel and M denotes number of EEG data from each channel.
		Output: FFT result with dimension N x M. N denotes number of channel and M denotes number of FFT data from each channel.
		"""
		data_fft = map(lambda x: np.fft.fft(x),all_channel_data)

		return data_fft

	def get_frequency(self,all_channel_data):
		"""
		Get frequency from computed fft for all channels.
		Input: Channel data with dimension N x M. N denotes number of channel and M denotes number of EEG data from each channel.
		Output: Frequency band from each channel: Delta, Theta, Alpha, Beta, and Gamma.
		"""
		#Length data channel
		L = len(all_channel_data[0])

		#Sampling frequency
		Fs = 128

		#Get fft data
		data_fft = self.do_fft(all_channel_data)

		#Compute frequency
		frequency = map(lambda x: abs(x//L),data_fft)
		frequency = map(lambda x: x[: L//2+1]*2,frequency)

		#creating 5 instances of frequency iterator
		f1,f2,f3,f4,f5=itertools.tee(frequency,5)

		#List frequency
		delta = np.array(list(map(lambda x: x[L*1//Fs-1: L*4//Fs],f1)))
		theta = np.array(list(map(lambda x: x[L*4//Fs-1: L*8//Fs],f2)))
		alpha = np.array(list(map(lambda x: x[L*5//Fs-1: L*13//Fs],f3)))
		beta =  np.array(list(map(lambda x: x[L*13//Fs-1: L*30//Fs],f4)))
		gamma = np.array(list(map(lambda x: x[L*30//Fs-1: L*50//Fs],f5)))

		return delta,theta,alpha,beta,gamma

	def get_feature(self,all_channel_data):
		"""
		Get feature from each frequency.
		Input: Channel data with dimension N x M. N denotes number of channel and M denotes number of EEG data from each channel.
		Output: Feature (standard deviasion and mean) from all frequency bands and channels with dimesion 1 x M (number of feature).
		"""

		(delta,theta,alpha,beta,gamma) = self.get_frequency(all_channel_data)
		delta_std = np.std(delta,axis=1)
		theta_std = np.std(theta,axis=1)
		alpha_std = np.std(alpha,axis=1)
		beta_std = np.std(beta,axis=1)
		gamma_std = np.std(gamma,axis=1)
		feature = np.array([delta_std,theta_std,alpha_std,beta_std,gamma_std])
		feature = feature.T
		feature = feature.ravel()

		return feature

	def predict_emotion(self,feature):
		"""
		Get arousal and valence class from feature.
		Input: Feature (standard deviasion and mean) from all frequency bands and channels with dimesion 1 x M (number of feature).
		Output: Class of emotion between 1 to 3 from each arousal and valence. 1 denotes low , 2 denotes neutral, and 3 denotes high .
		"""
		#Compute canberra with arousal training data
		distance_ar = list(map(lambda x:ss.distance.canberra(x,feature),self.train_arousal))
		#Compute canberra with valence training data
		distance_va = list(map(lambda x:ss.distance.canberra(x,feature),self.train_valence))

		#Compute 3 nearest index and distance value from arousal
		idx_nearest_ar = np.array(np.argsort(distance_ar)[:3])
		val_nearest_ar = np.array(np.sort(distance_ar)[:3])

		#Compute 3 nearest index and distance value from arousal
		idx_nearest_va = np.array(np.argsort(distance_va)[:3])
		val_nearest_va = np.array(np.sort(distance_va)[:3])

		#Compute comparation from first nearest and second nearest distance. If comparation less or equal than 0.7, then take class from the first nearest distance. Else take frequently class.
		#Arousal
		comp_ar = val_nearest_ar[0]/val_nearest_ar[1]
		if comp_ar<=0.7:
			result_ar = self.class_arousal[0,idx_nearest_ar[0]]
		else:
			result_ar = sst.mode(self.class_arousal[0,idx_nearest_ar])
			result_ar = float(result_ar[0])

		#Valence
		comp_va = val_nearest_va[0]/val_nearest_va[1]
		if comp_va<=0.7:
			result_va = self.class_valence[0,idx_nearest_va[0]]
		else:
			result_va = sst.mode(self.class_valence[0,idx_nearest_va])
			result_va = float(result_va[0])

		return result_ar,result_va

	def determine_emotion_class(self,feature):
		"""
		Get emotion class from feature.
		Input: Feature (standard deviasion) from all frequency bands and channels with dimesion 1 x M (number of feature).
		Output: Class of emotion between 1 to 5 according to Russel's Circumplex Model.
		"""
		class_ar,class_va = self.predict_emotion(feature)

		if class_ar==2.0 or class_va==2.0:
			emotion_class = 5
		elif class_ar==3.0 and class_va==1.0:
			emotion_class = 1
		elif class_ar==3.0 and class_va==3.0:
			emotion_class = 2
		elif class_ar==1.0 and class_va==3.0:
			emotion_class = 3
		elif class_ar==1.0 and class_va==1.0:
			emotion_class = 4
		return emotion_class

	def process_all_data(self,all_channel_data):
		"""
		Process all data from EEG data to predict emotion class.
		Input: Channel data with dimension N x M. N denotes number of channel and M denotes number of EEG data from each channel.
		Output: Class of emotion between 1 to 5 according to Russel's Circumplex Model. And send it to web ap
		"""
		#Get feature from EEG data
		feature = self.get_feature(all_channel_data)

		#Predict emotion class
		emotion_class = self.determine_emotion_class(feature)
		return emotion_class
#
	def send_result_to_window(self,emotion_class):
		"""
		Send emotion predict to web app.
		Input: Class of emotion between 1 to 5 according to Russel's Circumplex Model.
		Output: Send emotion prediction to web window.
		"""
		i1=cv2.imread('emoji/1.png')
		i2=cv2.imread('emoji/2.png')
		i3=cv2.imread('emoji/3.png')
		i4=cv2.imread('emoji/4.png')
		i5=cv2.imread('emoji/5.png')
		if emotion_class==1:
			cv2.imshow('image',i1)
		elif emotion_class==2:
			cv2.imshow('image',i2)
		elif emotion_class==3:
			cv2.imshow('image',i3)
		elif emotion_class==4:
			cv2.imshow('image',i4)
		else:
			cv2.imshow('image',i5)
		cv2.waitKey(0)
		cv2.destroyAllWindows()

	def main_process(self):

		"""
		Input: Get EEG data from DEAP, process all data (FFT, feature extraction, and classification), and predict the emotion.
		Output: Class of emotion between 1 to 5 according to Russel's Circumplex Model.
		"""
		#dataset NOT provided
		fname ="data/s02.dat"

		x = cPickle.load(open(fname, 'rb'), encoding="bytes")

		#feature vector formed from s01.dat and tested on s02.dat
		tr=10 #trial can be anywhere from (0,39)
		eeg_realtime=x[b'data'][tr]
		print(x[b'labels'][tr])
		eeg_raw=np.reshape(eeg_realtime,(40,8064))
		#slicing array as we only need first 32 channels
		eeg_raw=eeg_raw[:32,:]
		emotion_class=self.process_all_data(eeg_raw)
		print("Class of emotion="+str(emotion_class))
		if emotion_class==1:
			print("fear - nervous - stress - tense - upset")
		elif emotion_class==2:
			print("happy - alert - excited - elated")
		elif emotion_class==3:
			print("relax - calm - serene - contented")
		elif emotion_class==4:
			print("sad - depressed - lethargic - fatigue")
		else:
			print("Neutral")
		self.send_result_to_window(emotion_class)

#

if __name__ == "__main__":
	rte = predictEmotion()
	rte.main_process()


from kafka import KafkaConsumer
import os
import sys
import time
import socket
from tensorflow import keras
import numpy as np

print("We are here");
kafka_host = os.getenv('KAFKA_HOST_NAME')
if(kafka_host == None):
    print("Failed, no KAFKA_HOST_NAME environment variable was set")
    sys.exit(1)

if __name__ == "__main__":

    HOST, PORT = "trainer-model-service", 80
    data = "my data"
    # Create a socket (SOCK_STREAM means a TCP socket)
    while 1:
        try:
            # Connect to server and send data
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.connect((HOST, PORT))
        except socket.error:
            print("Couldn't connect to server, no model available");
            time.sleep(5);
            sock.close();
            continue
       sock.sendall(bytes(data + " ", "utf-8"))
        # Receive data from the server and shut down
        print("Receiving the model:");
        received = b'';
        numReceived = 0 ;
        while True:
            data = sock.recv(1048576);
            print("Recevied data, len = " + str(len(data)) + ", this is byte num " + str(numReceived) + "\n");
            numReceived+=1;
            if not data: break
            received += data;
        sock.close()
        break

    print("Len = " + str(len(received)));
    f = open('/model.h5', 'wb')
    f.write(received)
    f.close()
    print("Model receieved");
    model = keras.models.load_model("/model.h5")

    consumer  = KafkaConsumer("images", group_id="processor",  request_timeout_ms=120000,session_timeout_ms=100000, bootstrap_servers=kafka_host)

    for message in consumer:
       bytestream = bytes(message.value)
       newer = np.frombuffer(bytestream, dtype=np.uint8)
       newer = newer.reshape((28,28))
       newer = np.expand_dims(newer, axis=0)
       newer = np.expand_dims(newer, axis=3)
       scores = model.predict(newer);
       print("Prediction = " + str(scores));