import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
__all__ = ['SegmentationLosses', 'OhemCrossEntropy2d', 'OHEMSegmentationLosses']

class SegmentationLosses(nn.CrossEntropyLoss):
    """2D Cross Entropy Loss with Auxilary Loss"""
    def __init__(self, se_loss=False, se_weight=0.2, nclass=-1,
                 aux=False, aux_weight=0.4, weight=None,
                 ignore_index=-1):
        super(SegmentationLosses, self).__init__(weight, None, ignore_index)
        self.se_loss = se_loss
        self.aux = aux
        self.nclass = nclass
        self.se_weight = se_weight
        self.aux_weight = aux_weight
        self.bceloss = nn.BCELoss(weight)

    def forward(self, *inputs):
        if not self.se_loss and not self.aux:
            return super(SegmentationLosses, self).forward(*inputs)
        elif not self.se_loss:
            pred1, pred2, target = tuple(inputs)
            loss1 = super(SegmentationLosses, self).forward(pred1, target)
            loss2 = super(SegmentationLosses, self).forward(pred2, target)
            return loss1 + self.aux_weight * loss2
        elif not self.aux:
            pred, se_pred, target = tuple(inputs)
            se_target = self._get_batch_label_vector(target, nclass=self.nclass).type_as(pred)
            loss1 = super(SegmentationLosses, self).forward(pred, target)
            loss2 = self.bceloss(torch.sigmoid(se_pred), se_target)
            return loss1 + self.se_weight * loss2
        else:
            pred1, se_pred, pred2, target = tuple(inputs)
            se_target = self._get_batch_label_vector(target, nclass=self.nclass).type_as(pred1)
            loss1 = super(SegmentationLosses, self).forward(pred1, target)
            loss2 = super(SegmentationLosses, self).forward(pred2, target)
            loss3 = self.bceloss(torch.sigmoid(se_pred), se_target)
            return loss1 + self.aux_weight * loss2 + self.se_weight * loss3

    @staticmethod
    def _get_batch_label_vector(target, nclass):
        # target is a 3D Variable BxHxW, output is 2D BxnClass
        batch = target.size(0)
        tvect = Variable(torch.zeros(batch, nclass))
        for i in range(batch):
            hist = torch.histc(target[i].cpu().data.float(),
                               bins=nclass, min=0,
                               max=nclass-1)
            vect = hist>0
            tvect[i] = vect
        return tvect

# adapted from https://github.com/PkuRainBow/OCNet/blob/master/utils/loss.py
class OhemCrossEntropy2d(nn.Module):
    def __init__(self, ignore_label=-1, thresh=0.7, min_kept=100000, use_weight=True):
        super(OhemCrossEntropy2d, self).__init__()
        self.ignore_label = ignore_label
        self.thresh = float(thresh)
        self.min_kept = int(min_kept)
        if use_weight:
            print("w/ class balance")
            weight = torch.FloatTensor([0.8373, 0.918, 0.866, 1.0345, 1.0166, 0.9969, 0.9754,
                1.0489, 0.8786, 1.0023, 0.9539, 0.9843, 1.1116, 0.9037, 1.0865, 1.0955,
                1.0865, 1.1529, 1.0507])
            self.criterion = torch.nn.CrossEntropyLoss(weight=weight, ignore_index=ignore_label)
        else:
            print("w/o class balance")
            self.criterion = torch.nn.CrossEntropyLoss(ignore_index=ignore_label)

    def forward(self, predict, target, weight=None):
        """
        Args:
            predict:(n, c, h, w)
            target:(n, h, w)
            weight (Tensor, optional): a manual rescaling weight given to each class.
                                       If given, has to be a Tensor of size "nclasses"
        """
        assert not target.requires_grad
        assert predict.dim() == 4
        assert target.dim() == 3
        assert predict.size(0) == target.size(0), "{0} vs {1} ".format(predict.size(0), target.size(0))
        assert predict.size(2) == target.size(1), "{0} vs {1} ".format(predict.size(2), target.size(1))
        assert predict.size(3) == target.size(2), "{0} vs {1} ".format(predict.size(3), target.size(3))

        n, c, h, w = predict.size()
        input_label = target.data.cpu().numpy().ravel().astype(np.int32)
        x = np.rollaxis(predict.data.cpu().numpy(), 1).reshape((c, -1))
        input_prob = np.exp(x - x.max(axis=0).reshape((1, -1)))
        input_prob /= input_prob.sum(axis=0).reshape((1, -1))

        valid_flag = input_label != self.ignore_label
        valid_inds = np.where(valid_flag)[0]
        label = input_label[valid_flag]
        num_valid = valid_flag.sum()
        if self.min_kept >= num_valid:
            print('Labels: {}'.format(num_valid))
        elif num_valid > 0:
            prob = input_prob[:,valid_flag]
            pred = prob[label, np.arange(len(label), dtype=np.int32)]
            threshold = self.thresh
            if self.min_kept > 0:
                index = pred.argsort()
                threshold_index = index[ min(len(index), self.min_kept) - 1 ]
                if pred[threshold_index] > self.thresh:
                    threshold = pred[threshold_index]
            kept_flag = pred <= threshold
            valid_inds = valid_inds[kept_flag]

        label = input_label[valid_inds].copy()
        input_label.fill(self.ignore_label)
        input_label[valid_inds] = label
        valid_flag_new = input_label != self.ignore_label
        # print(np.sum(valid_flag_new))
        target = Variable(torch.from_numpy(input_label.reshape(target.size())).long().cuda())

        return self.criterion(predict, target)

class OHEMSegmentationLosses(OhemCrossEntropy2d):
    """2D Cross Entropy Loss with Auxilary Loss"""
    def __init__(self, se_loss=False, se_weight=0.2, nclass=-1,
                 aux=False, aux_weight=0.4, weight=None,
                 ignore_index=-1):
        super(OHEMSegmentationLosses, self).__init__(ignore_index)
        self.se_loss = se_loss
        self.aux = aux
        self.nclass = nclass
        self.se_weight = se_weight
        self.aux_weight = aux_weight
        self.bceloss = nn.BCELoss(weight)

    def forward(self, *inputs):
        if not self.se_loss and not self.aux:
            return super(OHEMSegmentationLosses, self).forward(*inputs)
        elif not self.se_loss:
            pred1, pred2, target = tuple(inputs)
            loss1 = super(OHEMSegmentationLosses, self).forward(pred1, target)
            loss2 = super(OHEMSegmentationLosses, self).forward(pred2, target)
            return loss1 + self.aux_weight * loss2
        elif not self.aux:
            pred, se_pred, target = tuple(inputs)
            se_target = self._get_batch_label_vector(target, nclass=self.nclass).type_as(pred)
            loss1 = super(OHEMSegmentationLosses, self).forward(pred, target)
            loss2 = self.bceloss(torch.sigmoid(se_pred), se_target)
            return loss1 + self.se_weight * loss2
        else:
            pred1, se_pred, pred2, target = tuple(inputs)
            se_target = self._get_batch_label_vector(target, nclass=self.nclass).type_as(pred1)
            loss1 = super(OHEMSegmentationLosses, self).forward(pred1, target)
            loss2 = super(OHEMSegmentationLosses, self).forward(pred2, target)
            loss3 = self.bceloss(torch.sigmoid(se_pred), se_target)
            return loss1 + self.aux_weight * loss2 + self.se_weight * loss3

    @staticmethod
    def _get_batch_label_vector(target, nclass):
        # target is a 3D Variable BxHxW, output is 2D BxnClass
        batch = target.size(0)
        tvect = Variable(torch.zeros(batch, nclass))
        for i in range(batch):
            hist = torch.histc(target[i].cpu().data.float(),
                               bins=nclass, min=0,
                               max=nclass-1)
            vect = hist>0
            tvect[i] = vect
        return tvect


import numpy as np
from sklearn.metrics import classification_report
from sklearn.model_selection import StratifiedKFold
from sklearn import metrics
import pandas as pd
import argparse

class Markov():

    def fasta2np(self, fasta):
        f = open(fasta)
        doc = f.readlines()
        list = []
        for seq in doc:
            if seq.startswith(">"):
                continue
            else:
                flag = seq.strip()
                a = [flag[i:i + 2] for i in range(len(flag) - 1)]
                list.append(a)
        m = np.array(list)
        return m

    '''m:训练集,p训练集中正例的个数'''

    def fit(self, m, p):
        pos = m[:p, :]

        neg = m[p:, :]

        '''pos的初始概率'''
        atgc1 = {}
        for i in range(len(pos)):
            if pos[i][0][0] in atgc1:

                atgc1[pos[i][0][0]] = atgc1[pos[i][0][0]] + 1
            else:
                atgc1[pos[i][0][0]] = 1
        for i, j in atgc1.items():
            atgc1[i] = atgc1[i] / len(pos)

        '''neg的初始概率'''
        atgc2 = {}
        for i in range(len(neg)):
            if neg[i][0][0] in atgc2:

                atgc2[neg[i][0][0]] = atgc2[neg[i][0][0]] + 1
            else:
                atgc2[neg[i][0][0]] = 1
        for i, j in atgc2.items():
            atgc2[i] = atgc2[i] / len(neg)

        '''pos的转移概率'''
        posdict = []
        for j in range(len(pos[0])):
            flag = {}
            a = 0
            t = 0
            g = 0
            c = 0
            for i in range(len(pos)):
                if pos[i][j][0] == "A":
                    a += 1
                elif pos[i][j][0] == "T":
                    t += 1
                elif pos[i][j][0] == "G":
                    g += 1
                else:
                    c += 1

                if pos[i, j] in flag:
                    flag[pos[i, j]] = flag[pos[i, j]] + 1
                else:
                    flag[pos[i, j]] = 1

            for i in flag:

                if i[0] == "A":
                    flag[i] = flag[i] / a
                elif i[0] == "T":
                    flag[i] = flag[i] / t
                elif i[0] == "G":
                    flag[i] = flag[i] / g
                else:
                    flag[i] = flag[i] / c

            posdict.append(flag)

        '''neg的转移概率'''
        negdict = []
        for j in range(len(neg[0])):
            flag = {}
            a = 0
            t = 0
            g = 0
            c = 0
            for i in range(len(neg)):

                if neg[i][j][0] == "A":
                    a += 1
                elif neg[i][j][0] == "T":
                    t += 1
                elif neg[i][j][0] == "G":
                    g += 1
                else:
                    c += 1

                if neg[i, j] in flag:
                    flag[neg[i, j]] = flag[neg[i, j]] + 1
                else:
                    flag[neg[i, j]] = 1

            for i in flag:

                if i[0] == "A":
                    flag[i] = flag[i] / a
                elif i[0] == "T":
                    flag[i] = flag[i] / t
                elif i[0] == "G":
                    flag[i] = flag[i] / g
                else:
                    flag[i] = flag[i] / c

            negdict.append(flag)

        posdict = np.array(posdict)
        negdict = np.array(negdict)

        return posdict, negdict, atgc1, atgc2

    def predict(self, test_set, posdict, negdict, atgc1, atgc2):
        posfea = []
        for i in range(len(test_set)):
            pos = 1
            if test_set[i][0][0] in atgc1:
                pos = pos * atgc1[test_set[i][0][0]]
            for j in range(len(test_set[i])):
                if test_set[i][j] in posdict[j]:
                    pos = pos * (posdict[j][test_set[i][j]])
                else:
                    pos = pos * 1
            posfea.append(pos)

        negfea = []
        for i in range(len(test_set)):
            neg = 1
            if test_set[i][0][0] in atgc1:
                neg = neg * atgc2[test_set[i][0][0]]
            for j in range(len(test_set[i])):
                if test_set[i][j] in negdict[j]:
                    neg = neg * (negdict[j][test_set[i][j]])
                else:
                    neg = neg * 1
            negfea.append(neg)

        label = np.array(posfea) / np.array(negfea)

        return label

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-trainfasta', required=True, help="trainfasta file name")
    parser.add_argument('-trainpos', required=True, type=int, help="trainpos")
    parser.add_argument('-trainneg', required=True, type=int, help="trainneg")
    parser.add_argument('-testfasta',  help="testfasta file name")
    parser.add_argument('-testpos',  type=int, help="testpos")
    parser.add_argument('-testneg',  type=int, help="testneg")
    parser.add_argument('-cv', default=10, type=int, help="cv")

    args = parser.parse_args()
    print(args)

    m = Markov()
    x = m.fasta2np(args.trainfasta)
    y = np.array([0] * args.trainpos + [1] * args.trainneg)

    if args.testfasta:
        x1 = m.fasta2np(args.testfasta)
        y1 = np.array([0] * args.testpos + [1] * args.testneg)

    if args.testfasta:
        posdict, negdict, atgc1, atgc2 = m.fit(x, args.trainpos)

        label = m.predict(x1, posdict, negdict, atgc1, atgc2)
        l = []
        for i in label:
            if i < 1:
                l.append(1)
            else:
                l.append(0)
        l=np.array(l)

        print("ACC:{}".format(metrics.accuracy_score(y1,l)))
        print("MCC:{}\n".format(metrics.matthews_corrcoef(y1,l)))
        print(classification_report(y1, l))
        print("confusion matrix\n")
        print(pd.crosstab(pd.Series(y1, name='Actual'), pd.Series(l, name='Predicted')))
    else:
        skf = StratifiedKFold(n_splits=args.cv)
        acc = []
        testlabel=[]
        truelabel=[]
        for train_index, test_index in skf.split(x, y):

            X_train, X_test = x[train_index], x[test_index]
            y_train, y_test = y[train_index], y[test_index]

            num = 0
            for i in y_train:
                if i == 0:
                    num += 1

            posdict, negdict, atgc1, atgc2 = m.fit(X_train, num)
            label = m.predict(X_test, posdict, negdict, atgc1, atgc2)

            l = []
            for i in label:
                if i < 1:
                    l.append(1)
                else:
                    l.append(0)

            acc.append(metrics.accuracy_score(y_test, l))
            testlabel+=l
            truelabel+=y_test.tolist()

            # print("ACC:{}".format(metrics.accuracy_score(y_test, l)))
        # print(np.mean(acc))
        print("ACC:{}".format(metrics.accuracy_score(truelabel,testlabel)))
        print("MCC:{}\n".format(metrics.matthews_corrcoef(truelabel, testlabel)))
        print(classification_report(truelabel, testlabel))
        print("confusion matrix\n")
        print(pd.crosstab(pd.Series(truelabel, name='Actual'), pd.Series(testlabel, name='Predicted')))

if __name__ == '__main__':
    main()


import copy
import hashlib
import pickle
from datetime import datetime, timezone

import numpy as np
import pandas as pd
import sqlalchemy.sql
from sqlalchemy.exc import SQLAlchemyError

from . import dfdict

def serialize(data):
    """Eventually turn some elements into their binary representation.

    Parameters
    ----------
    data : dict
        raw data

    Returns
    -------
    ndata : list
        improved data

    """
    blobbed_types = [list, dict, dfdict.DFdict, np.ndarray]
    sorted_data = dict(sorted(data.items()))
    ndata = []
    for d in sorted_data.values():
        if type(d) in blobbed_types:
            ndata.append(pickle.dumps(d))
        else:
            ndata.append(d)
    return tuple(ndata)

def deserialize(data):
    """Undo the binary representation.

    Parameters
    ----------
    data : list
        raw data

    Returns
    -------
    obj : dict
        composed object

    """
    obj = {}
    for f, el in data.__dict__.items():
        if f[0] == "_":
            continue
        if isinstance(el, bytes):
            obj[f] = pickle.loads(el)
            if isinstance(obj[f], dict) and "__msgs__" in obj[f]:
                obj[f] = dfdict.DFdict.from_document(obj[f])
        else:
            obj[f] = el
    return obj

def add_hash(record):
    """Add a hash value as last element to the record.

    Parameters
    ----------
    record : tuple
        data

    Returns
    -------
    hashed_record : tuple
        data + hash

    """
    h = hashlib.sha256(pickle.dumps(record))
    return (*record, h.digest().hex())

def prepare_records(base, updates):
    """Generate all records from the base.

    Parameters
    ----------
    base : dict
        base record
    updates : list(dict)
        update directives

    Returns
    -------
    documents : list(dict)
        list of all dictionaries
    df : pandas.DataFrame
        all records ready for insertion

    """
    records = []
    documents = []
    for upd in updates:
        document = copy.deepcopy(base)
        document.update(upd)
        documents.append(document)
        # add the hash when there is NO meta data around
        serialized_record = serialize(document)
        hashed_record = add_hash(serialized_record)
        document["hash"] = hashed_record[-1]
        records.append(hashed_record)
    return (documents, pd.DataFrame(records, columns=(list(base.keys()) + ["hash"])))

def insertmany(session, table, df):
    """Insert all records into the DB.

    Parameters
    ----------
    session : sqlalchemy.orm.session.Session
        database
    table : sqlalchemy.schema.Table
        target table
    df : pandas.DataFrame
        dataframe all records

    """
    try:
        session.bulk_insert_mappings(table, df.to_dict(orient="records"))
        # TODO: do we want to commit here or somewhere else?
        session.commit()
    except SQLAlchemyError:
        session.rollback()

def insertnew(session, table, df):
    """Insert all records that do not exist yet (determined by hash).

    Parameters
    ----------
    session : sqlalchemy.orm.session.Session
        database
    table : sqlalchemy.schema.Table
        target table
    df : pandas.DataFrame
        dataframe all records

    """
    # TODO: why the hash field is a 1-tuple?
    hashes = [h[0] for h in session.query(table.hash).all()]
    new_records = df[~df["hash"].isin(hashes)]
    insertmany(session, table, new_records)

class RetrieveError(KeyError):
    objects = None

    def __init__(self, descr, objs=None):
        if objs is not None:
            msg = "the following {self.objects} have been found:"
            msg += "\n- " + "\n- ".join(objs) + "\n"
            print(msg)
        super().__init__(descr)

class HashError(RetrieveError):
    objects = "hashes"

class UIDError(RetrieveError):
    objects = "pids"

def select_by_hash(session, table_object, hash_partial):
    """Find a record by its partial hash.

    The hash provided is considered to be the first ``len(hash_partial)``
    carachters of a full hash in the table.

    Parameters
    ----------
    session : sqlalchemy.orm.session.Session
        DB ORM session
    table_object : sqlalchemy.schema.Table
        table object
    hash_partial : str
        hash identifier

    Returns
    -------
    dict
        record

    """
    available = (
        session.query(table_object)
        .filter(
            sqlalchemy.sql.func.substr(table_object.hash, 1, len(hash_partial))
            == hash_partial
        )
        .all()
    )
    # too much?
    if len(available) > 1:
        raise HashError("hash is not unique", [a.hash for a in available])
    if len(available) < 1:
        raise HashError("hash not found")
    # deserialize the thing
    return deserialize(available[0])

def select_by_uid(session, table_object, uid):
    """Find a record by its uid (unique identifier).

    Parameters
    ----------
    session : sqlalchemy.orm.session.Session
        DB ORM session
    table_object : sqlalchemy.schema.Table
        table object
    hash_partial : str
        hash identifier

    Returns
    -------
    dict
        record

    """
    available = session.query(table_object).filter(table_object.uid == uid).all()
    # too much?
    if len(available) > 1:
        raise UIDError("uid is not unique", [a.hash for a in available])
    if len(available) < 1:
        raise UIDError("uid not found")
    # deserialize the thing
    return deserialize(available[0])

def select_by_position(session, table_object, pos):
    """Find a record by its position.

    Negative values are supported, and used as positions from the end (exact
    same logic of python :class:`list`).

    Note
    ----
    This is more expensive than using :func:`select_uid`, since the whole table
    is retrieved, and only after a single member is selected.

    Parameters
    ----------
    session : sqlalchemy.orm.session.Session
        DB ORM session
    table_object : sqlalchemy.schema.Table
        table object
    hash_partial : str
        hash identifier

    Returns
    -------
    dict
        record

    """
    return deserialize(session.query(table_object).all()[pos])

def select_all(session, table_object):
    """Collect all records.

    Parameters
    ----------
    session : sqlalchemy.orm.session.Session
        DB ORM session
    table_object : sqlalchemy.schema.Table
        table object

    Returns
    -------
    list(dict)
        list of records

    """
    available = session.query(table_object).all()
    return [deserialize(a) for a in available]

def update_atime(session, table_object, uids):
    """Update rows access time to now.

    Parameters
    ----------
    session : sqlalchemy.orm.session.Session
        DB ORM session
    table_object : sqlalchemy.schema.Table
        table object
    uids : list(int)
        unique identifiers of rows to update

    """
    try:
        for row in session.query(table_object).filter(table_object.uid.in_(uids)):
            row.atime = datetime.now(timezone.utc)
        session.commit()
    except SQLAlchemyError:
        session.rollback()

def truncate(session, table_object):
    """Empty table.

    Parameters
    ----------
    session : sqlalchemy.orm.session.Session
        DB ORM session
    table_object : sqlalchemy.schema.Table
        table object

    """
    try:
        session.query(table_object).delete()
        session.commit()
    except SQLAlchemyError:
        session.rollback()

def remove(session, table_object, uids):
    """Remove given rows from chosen table.

    Parameters
    ----------
    session : sqlalchemy.orm.session.Session
        DB ORM session
    table_object : sqlalchemy.schema.Table
        table object
    uids : list(int)
        unique identifiers of rows to delete

    """
    try:
        session.query(table_object).filter(table_object.uid.in_(uids)).delete()
        session.commit()
    except SQLAlchemyError:
        session.rollback()


from models.Discriminator import Discriminator
from models.VGG19 import Vgg19
from models.DINet import DINet
from models.Syncnet import SyncNetPerception
from models.color_syncnet import SyncNet_color  as SyncNet
from utils.training_utils import get_scheduler, update_learning_rate,GANLoss
from config.config import DINetTrainingOptions
from sync_batchnorm import convert_model
from torch.utils.data import DataLoader
from dataset.dataset_DINet_clip_colorsync import DINetDataset

import random
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import os
import torch.nn.functional as F

from torch.utils.tensorboard import SummaryWriter
import cv2

def _load(checkpoint_path):
    use_cuda = torch.cuda.is_available()
    if use_cuda:
        checkpoint = torch.load(checkpoint_path)
    else:
        checkpoint = torch.load(checkpoint_path,
                                map_location=lambda storage, loc: storage)
    return checkpoint

def load_checkpoint(path, model):
    print("Load checkpoint from: {}".format(path))
    checkpoint = _load(path)
    s = checkpoint["state_dict"]
    new_s = {}
    for k, v in s.items():
        new_s[k.replace('module.', '')] = v
    model.load_state_dict(new_s)
    return model

logloss = nn.BCELoss()
def cosine_loss(a, v, y):
    d = nn.functional.cosine_similarity(a, v)
    loss = logloss(d.unsqueeze(1), y)
    return loss

if __name__ == "__main__":
    '''
            clip training code of DINet
            in the resolution you want, using clip training code after frame training

        '''
    # load config
    opt = DINetTrainingOptions().parse_args()

    writer = SummaryWriter('./logs/'  + opt.result_path.split('/')[-1])

    random.seed(opt.seed)
    np.random.seed(opt.seed)
    torch.cuda.manual_seed(opt.seed)
    # load training data
    train_data = DINetDataset(opt.train_data,opt.augment_num,opt.mouth_region_size)
    training_data_loader = DataLoader(dataset=train_data,  batch_size=opt.batch_size, shuffle=True,drop_last=True)
    train_data_length = len(training_data_loader)
    # init network
    opt.audio_channel = 256 #######
    net_g = DINet(opt.source_channel,opt.ref_channel,opt.audio_channel).cuda()
    net_dI = Discriminator(opt.source_channel ,opt.D_block_expansion, opt.D_num_blocks, opt.D_max_features).cuda()
    net_dV = Discriminator(opt.source_channel * 5, opt.D_block_expansion, opt.D_num_blocks, opt.D_max_features).cuda()
    net_vgg = Vgg19().cuda()

    # import pdb
    # pdb.set_trace()
    # for sync net
    # net_lipsync = SyncNetPerception(opt.pretrained_syncnet_path).cuda()
    net_lipsync = SyncNet().cuda()
    for p in net_lipsync.parameters():
        p.requires_grad = False
    # pdb.set_trace()
    net_lipsync = load_checkpoint(opt.pretrained_syncnet_path, net_lipsync)
    net_lipsync = net_lipsync.cuda()

    # parallel
    net_g = nn.DataParallel(net_g)
    net_g = convert_model(net_g)
    net_dI = nn.DataParallel(net_dI)
    net_dV = nn.DataParallel(net_dV)
    net_vgg = nn.DataParallel(net_vgg)
    # setup optimizer
    optimizer_g = optim.Adam(net_g.parameters(), lr=opt.lr_g)
    optimizer_dI = optim.Adam(net_dI.parameters(), lr=opt.lr_dI)
    optimizer_dV = optim.Adam(net_dV.parameters(), lr=opt.lr_dI)
    ## load frame trained DInet weight
    print('loading frame trained DINet weight from: {}'.format(opt.pretrained_frame_DINet_path))
    checkpoint = torch.load(opt.pretrained_frame_DINet_path)
    net_g.load_state_dict(checkpoint['state_dict']['net_g'])
    # set criterion
    criterionGAN = GANLoss().cuda()
    criterionL1 = nn.L1Loss().cuda()
    criterionMSE = nn.MSELoss().cuda()
    # set scheduler
    net_g_scheduler = get_scheduler(optimizer_g, opt.non_decay, opt.decay)
    net_dI_scheduler = get_scheduler(optimizer_dI, opt.non_decay, opt.decay)
    net_dV_scheduler = get_scheduler(optimizer_dV, opt.non_decay, opt.decay)
    # set label of syncnet perception loss
    real_tensor = torch.tensor(1.0).cuda()
    # start train
    for epoch in range(opt.start_epoch, opt.non_decay+opt.decay+1):
        net_g.train()
        for iteration, data in enumerate(training_data_loader):
            # forward
            source_clip,source_clip_mask, reference_clip,deep_speech_clip,deep_speech_full, sync_mel, sync_y = data

            source_clip = torch.cat(torch.split(source_clip, 1, dim=1), 0).squeeze(1).float().cuda()
            source_clip_mask = torch.cat(torch.split(source_clip_mask, 1, dim=1), 0).squeeze(1).float().cuda()
            reference_clip = torch.cat(torch.split(reference_clip, 1, dim=1), 0).squeeze(1).float().cuda()
            deep_speech_clip = torch.cat(torch.split(deep_speech_clip, 1, dim=1), 0).squeeze(1).float().cuda()
            deep_speech_full = deep_speech_full.float().cuda()

            fake_out = net_g(source_clip_mask,reference_clip,deep_speech_clip)
            fake_out_half = F.avg_pool2d(fake_out, 3, 2, 1, count_include_pad=False)
            source_clip_half = F.interpolate(source_clip, scale_factor=0.5, mode='bilinear')

            # crop sync input fake data
            fake_out_tmp = fake_out.clone()
            _, _, h_fake, w_fake = fake_out_tmp.shape
            fake_out_tmp = fake_out_tmp[:, :, int(h_fake * 0.380) : int(h_fake * 0.818), :]
            fake_out_tmp = torch.split(fake_out_tmp, opt.batch_size, dim=0)
            sync_fake_input = [[] for i in range(opt.batch_size)]
            for ref_id in range(5):
                for batch_id in range(opt.batch_size):
                    tmp_img = fake_out_tmp[ref_id][batch_id, :, :, :].permute(1, 2, 0).detach().cpu().numpy()[:, :, ::-1]
                    tmp_img = cv2.resize(tmp_img, (96, 48))
                    sync_fake_input[batch_id].append(tmp_img)
            batch_all_person_data = []
            for batch_single_person_data in sync_fake_input:
                batch_single_person_array = np.array(batch_single_person_data)
                batch_single_person_array = np.concatenate(batch_single_person_array, axis=2)
                batch_single_person_array = batch_single_person_array.transpose(2, 0, 1)
                batch_all_person_data.append(batch_single_person_array)
            sync_fake_img_tensor = torch.FloatTensor(batch_all_person_data)

            # prepare syncnet input data
            sync_mel = sync_mel.cuda()
            sync_fake_img_tensor = sync_fake_img_tensor.cuda()
            sync_y = sync_y.cuda()

            # vis network output img
            if iteration % 100 == 0:
                fake_frame = fake_out[0, :, :, :].squeeze(0).permute(1, 2, 0).detach().cpu().numpy() * 255
                real_frame = source_clip[0, :, :, :].permute(1, 2, 0).detach().cpu().numpy() * 255
                fake_img_name = "./vis_img_clip/epoch" + str(epoch).zfill(4) + '_' + str(iteration).zfill(5) + '_fake.jpg'
                real_img_name = "./vis_img_clip/epoch" + str(epoch).zfill(4) + '_' + str(iteration).zfill(5) + '_real.jpg'
                cv2.imwrite(fake_img_name, fake_frame[:, :, ::-1])
                cv2.imwrite(real_img_name, real_frame[:, :, ::-1])

            # (1) Update DI network
            optimizer_dI.zero_grad()
            _,pred_fake_dI = net_dI(fake_out)
            loss_dI_fake = criterionGAN(pred_fake_dI, False)
            _,pred_real_dI = net_dI(source_clip)
            loss_dI_real = criterionGAN(pred_real_dI, True)
            # Combined DI loss
            # loss_dI = (loss_dI_fake + loss_dI_real) * 0.5               # GAN Loss D(true + Fake)  using  Discriminator-image-model
            loss_dI = (loss_dI_fake + loss_dI_real) * 0.5 * 1.5              # GAN Loss D(true + Fake)  using  Discriminator-image-model
            loss_dI.backward(retain_graph=True)
            optimizer_dI.step()

            # (2) Update DV network
            optimizer_dV.zero_grad()
            condition_fake_dV = torch.cat(torch.split(fake_out, opt.batch_size, dim=0), 1)
            _, pred_fake_dV = net_dV(condition_fake_dV)
            loss_dV_fake = criterionGAN(pred_fake_dV, False)
            condition_real_dV = torch.cat(torch.split(source_clip, opt.batch_size, dim=0), 1)
            _, pred_real_dV = net_dV(condition_real_dV)
            loss_dV_real = criterionGAN(pred_real_dV, True)
            # Combined DV loss
            loss_dV = (loss_dV_fake + loss_dV_real) * 0.5 * 1.5              # GAN Loss D(true + Fake) using  Discriminator-video-model
            # loss_dV = (loss_dV_fake + loss_dV_real) * 0.5              # GAN Loss D(true + Fake) using  Discriminator-video-model
            loss_dV.backward(retain_graph=True)
            optimizer_dV.step()

            # (2) Update DINet
            optimizer_g.zero_grad()
            # compute perception loss
            perception_real = net_vgg(source_clip)
            perception_fake = net_vgg(fake_out)
            perception_real_half = net_vgg(source_clip_half)
            perception_fake_half = net_vgg(fake_out_half)
            loss_g_perception = 0
            for i in range(len(perception_real)):
                loss_g_perception += criterionL1(perception_fake[i], perception_real[i])
                loss_g_perception += criterionL1(perception_fake_half[i], perception_real_half[i])
            loss_g_perception = (loss_g_perception / (len(perception_real) * 2)) * opt.lamb_perception        # perception-Loss  using  DINet(net_g)

            _, pred_fake_dI = net_dI(fake_out)
            _, pred_fake_dV = net_dV(condition_fake_dV)
            # # gan dI loss
            loss_g_dI = criterionGAN(pred_fake_dI, True)                 # GAN Loss G  using Discriminator-image-model
            # # gan dV loss
            loss_g_dV = criterionGAN(pred_fake_dV, True)                 # GAN Loss G  using Discriminator-video-model

            # ## color sync loss
            sync_a, sync_v = net_lipsync(sync_mel, sync_fake_img_tensor)
            loss_sync = cosine_loss(sync_a, sync_v, sync_y)

            # combine all losses
            loss_g =   loss_g_perception +  1.5 * (loss_g_dI +loss_g_dV) + loss_sync * 1.0
            # loss_g =   loss_g_perception +  loss_g_dI +loss_g_dV + loss_sync
            loss_g.backward()
            optimizer_g.step()

            print(
                "===> Epoch[{}]({}/{}):  Loss_Gan_D_image: {:.4f} Loss_Gan_G_image: {:.4f} Loss_Gan_D_video: {:.4f} Loss_Gan_G_video: {:.4f} Loss_perception: {:.4f} Loss_sync: {:.4f} lr_g = {:.7f} ".format(
                    epoch, iteration, len(training_data_loader), float(loss_dI), float(loss_g_dI),float(loss_dV), float(loss_g_dV), float(loss_g_perception),float(loss_sync),
                    optimizer_g.param_groups[0]['lr']))

            writer.add_scalar("Loss_Gan_D_image", float(loss_dI), epoch*len(training_data_loader)+iteration)
            writer.add_scalar("Loss_Gan_G_image", float(loss_g_dI), epoch*len(training_data_loader)+iteration)
            writer.add_scalar("Loss_Gan_D_video", float(loss_dV), epoch*len(training_data_loader)+iteration)
            writer.add_scalar("Loss_Gan_G_video", float(loss_g_dV), epoch*len(training_data_loader)+iteration)
            writer.add_scalar("Loss_perception", float(loss_g_perception), epoch*len(training_data_loader)+iteration)
            writer.add_scalar("Loss_sync", float(loss_sync), epoch*len(training_data_loader)+iteration)
            writer.add_scalar("lr_g", float(optimizer_g.param_groups[0]['lr']), epoch*len(training_data_loader)+iteration)

        update_learning_rate(net_g_scheduler, optimizer_g)
        update_learning_rate(net_dI_scheduler, optimizer_dI)
        update_learning_rate(net_dV_scheduler, optimizer_dV)

        writer.add_scalar("Loss_DI_epoch", float(loss_dI), epoch)
        writer.add_scalar("Loss_GI_epoch", float(loss_g_dI), epoch)
        writer.add_scalar("Loss_DV_epoch", float(loss_dV), epoch)
        writer.add_scalar("Loss_Gan_G_video", float(loss_g_dV), epoch)
        writer.add_scalar("Loss_perception_epoch", float(loss_g_perception), epoch)
        writer.add_scalar("Loss_sync_epoch", float(loss_sync), epoch)
        writer.add_scalar("lr_g_epoch", float(optimizer_g.param_groups[0]['lr']), epoch)

        # checkpoint
        if epoch %  opt.checkpoint == 0:
            if not os.path.exists(opt.result_path):
                os.mkdir(opt.result_path)
            model_out_path = os.path.join(opt.result_path, 'netG_model_epoch_{}.pth'.format(epoch))
            states = {
                'epoch': epoch + 1,
                'state_dict': {'net_g': net_g.state_dict(),'net_dI': net_dI.state_dict(),'net_dV': net_dV.state_dict()},
                'optimizer': {'net_g': optimizer_g.state_dict(), 'net_dI': optimizer_dI.state_dict(), 'net_dV': optimizer_dV.state_dict()}
            }
            torch.save(states, model_out_path)
            print("Checkpoint saved to {}".format(epoch))


import streamlit as st
import requests
import pandas as pd
import plotly.express as px

st.set_page_config(layout = 'wide')

def formata_numero(valor, prefixo = ''):
    for unidade in ['', 'mil']:
        if valor < 1000:
            return f'{prefixo} {valor:.2f} {unidade}'
        valor /= 1000
    return f'{prefixo} {valor:.2f} milhões'

#TITULO
st.title('DASHBOARD DE VENDAS:shopping_trolley:')

#Coleta dos Dados
url = 'https://labdados.com/produtos'

regioes = ['Brasil', 'Centro-Oeste', 'Nordeste', 'Norte', 'Sudeste', 'Sul']

st.sidebar.title('Filtros')

regiao = st.sidebar.selectbox('Região', regioes)

if regiao == 'Brasil':
    regiao = ''

todos_anos = st.sidebar.checkbox('Dados de todo o período', value= True)

if todos_anos:
    ano = ''
else:
    ano =  st.sidebar.slider('Ano', 2020, 2023)

query_string = {'regiao':regiao.lower(), 'ano':ano}
response = requests.get(url, params=query_string)

dados = pd.DataFrame.from_dict(response.json())

dados['Data da Compra'] = pd.to_datetime(dados['Data da Compra'], format='%d/%m/%Y')

filtro_vendedores = st.sidebar.multiselect('Vendedores', dados['Vendedor'].unique())
if filtro_vendedores:
    dados = dados[dados['Vendedor'].isin(filtro_vendedores)]

##TABELAS

#Tabelas receita
receitas_estados = dados.groupby('Local da compra')[['Preço']].sum()
receitas_estados = dados.drop_duplicates(subset='Local da compra')[['Local da compra', 'lat', 'lon']].merge(receitas_estados, left_on = 'Local da compra', right_index = True).sort_values('Preço', ascending=False)

receita_mensal = dados.set_index('Data da Compra').groupby(pd.Grouper(freq= 'M'))['Preço'].sum().reset_index()
receita_mensal['Ano'] = receita_mensal['Data da Compra'].dt.year
receita_mensal['Mês'] = receita_mensal['Data da Compra'].dt.month_name()

receita_categorias = dados.groupby('Categoria do Produto')[['Preço']].sum().sort_values('Preço', ascending= False)

#Tabela quantidade de vendas
# Calculando a contagem de vendas por 'Local da compra'
quantidade_venda = dados.groupby('Local da compra').size().reset_index(name='Quantidade de vendas')

# Obtendo os dados de latitude e longitude
dados_lat_lon = dados.drop_duplicates(subset='Local da compra')[['Local da compra', 'lat', 'lon']]

# Mesclando os dois DataFrames com concatenação
quantidade_venda = pd.concat([dados_lat_lon.set_index('Local da compra'), quantidade_venda.set_index('Local da compra')], axis=1).reset_index()

# Ordenando pela quantidade de vendas em ordem decrescente
quantidade_venda = quantidade_venda.sort_values('Quantidade de vendas', ascending=False)

# Vendas Mensal

# Supondo que a coluna 'Data da Compra' contenha a data da compra no formato apropriado (você pode precisar ajustar isso)
dados['Data da Compra'] = pd.to_datetime(dados['Data da Compra'])

# Extrair o mês da coluna 'Data da Compra' e converter para string
dados['Mês'] = dados['Data da Compra'].dt.to_period('M').astype(str)

# Agrupar por Mês e Local da compra e contar as vendas
vendas_por_mes_local = dados.groupby(['Mês', 'Local da compra']).size().reset_index(name='Quantidade de vendas')

vendas_categoria = dados.groupby('Categoria do Produto')[['Local da compra']].count().sort_values('Local da compra', ascending= False)

#Tabela Vendedores

vendedores = pd.DataFrame(dados.groupby('Vendedor')['Preço'].agg(['sum', 'count']))

##GRAFICOS

#graficos receita
fig_mapa_receita = px.scatter_geo(receitas_estados,
                                  lat='lat',
                                  lon='lon',
                                  scope='south america',
                                  size='Preço',
                                  template='seaborn',
                                  hover_name='Local da compra',
                                  title='Receita por estado')

fig_receita_mensal = px.line(receita_mensal,
                             x = 'Mês',
                             y = 'Preço',
                             markers= True,
                             range_y= (0, receita_mensal.max()),
                             color= 'Ano',
                             title= 'Receita Mensal')

fig_receita_mensal.update_layout(yaxis_title = 'Receita')

fig_receita_estado = px.bar(receitas_estados.head(),
                            x = 'Local da compra',
                            y = 'Preço',
                            color = 'Local da compra',
                            text_auto= True,
                            title= 'Top estados (receita)')

fig_receita_estado.update_layout(yaxis_title = 'Receita')

fig_receita_categorias = px.bar(receita_categorias,
                                text_auto= True,
                                title= 'Receita por categoria')

fig_receita_categorias.update_layout(yaxis_title = 'Quantidade de Vendas')

#graficos quantidade de venda

fig_mapa_venda = px.scatter_geo(quantidade_venda,
                                lat='lat',  # Substitua 'lat' pelo nome correto da coluna de latitude
                                lon='lon',  # Substitua 'lon' pelo nome correto da coluna de longitude
                                scope='south america',
                                size='Quantidade de vendas',  # Certifique-se de que 'Preço' é a coluna correta para definir o tamanho
                                template='seaborn',
                                hover_name='Local da compra',
                                title='Quantidade venda por estado')

fig_vendas_mensal = px.line(vendas_por_mes_local, x='Mês', y='Quantidade de vendas', color='Local da compra',
              labels={'Mês': 'Mês', 'Quantidade de vendas': 'Quantidade de Vendas'},
              title='Quantidade de Vendas por Mês')

fig_vendas_mensal.update_layout(yaxis_title = 'Quantidade de Vendas')

fig_vendas_estado = px.bar(quantidade_venda.head(),
                            x = 'Local da compra',
                            y = 'Quantidade de vendas',
                            color = 'Local da compra',
                            text_auto= True,
                            title= 'Top estados (vendas)')

fig_receita_estado.update_layout(yaxis_title = 'Receita')

fig_vendas_categorias = px.bar(vendas_categoria,
                                text_auto= True,
                                title= 'Vendas Categorias(Estado)')

fig_receita_categorias.update_layout(yaxis_title = 'Quantidade de Vendas')

##VISUALIZACAO NO STREAMLITE

#Abas

aba1, aba2, aba3 = st.tabs(['Receitas', 'Quantidade de Vendas', 'Vendedores'])

#Metricas

with aba1:

    col1, col2 = st.columns(2)
    with col1:
        st.metric('Receita',formata_numero(dados['Preço'].sum(), 'R$'))
        st.plotly_chart(fig_mapa_receita,use_container_width=True)
        st.plotly_chart(fig_receita_estado,use_container_width=True)
    with col2:
        st.metric('Quantidade de Vendas', formata_numero(dados.shape[0]))
        st.plotly_chart(fig_receita_mensal, use_container_width=True)
        st.plotly_chart(fig_receita_categorias, use_container_width=True)

with aba2:

    col1, col2 = st.columns(2)
    with col1:
        st.metric('Receita',formata_numero(dados['Preço'].sum(), 'R$'))
        st.plotly_chart(fig_mapa_venda, use_container_width= True)
        st.plotly_chart(fig_vendas_estado, use_container_width=True)
    with col2:
        st.metric('Quantidade de Vendas', formata_numero(dados.shape[0]))
        st.plotly_chart(fig_vendas_mensal, use_container_width=True)
        st.plotly_chart(fig_vendas_categorias, use_container_width=True)

with aba3:
    qtd_vendedores = st.number_input('Quantidades de Vendedores', 2 , 10 ,5)

    col1, col2 = st.columns(2)
    with col1:
        st.metric('Receita',formata_numero(dados['Preço'].sum(), 'R$'))
        fig_receita_vendedores = px.bar(vendedores[['sum']].sort_values('sum', ascending= False).head(qtd_vendedores),
                                        x = 'sum',
                                        y = vendedores[['sum']].sort_values('sum', ascending= False).head(qtd_vendedores).index,
                                        text_auto= True,
                                        color = 'sum',
                                        title= f'Top {qtd_vendedores} vendedores (receita)')

        st.plotly_chart(fig_receita_vendedores, use_container_width=True)
    with col2:
        st.metric('Quantidade de Vendas', formata_numero(dados.shape[0]))
        fig_vendas_vendedores = px.bar(vendedores[['count']].sort_values('count', ascending= False).head(qtd_vendedores),
                                        x = 'count',
                                        y = vendedores[['count']].sort_values('count', ascending= False).head(qtd_vendedores).index,
                                        text_auto= True,
                                        color = 'count',
                                        title= f'Top {qtd_vendedores} vendedores (vendas)')

        st.plotly_chart(fig_vendas_vendedores, use_container_width=True)


import numpy as np
import cv2

cap = cv2.VideoCapture(0)

# 设置摄像头分辨率
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 2560)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

fourcc = cv2.VideoWriter_fourcc(*'mp4v')
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
out = cv2.VideoWriter('/home/eaibot71/test1/photo_video/output.mp4',fourcc, 5.0, (width,height))

while(cap.isOpened()):
    ret, frame = cap.read()
    if ret==True:
        frame = cv2.flip(frame,1)
#        out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))   将视频转换为灰色的源
        out.write(frame)

        cv2.imshow('frame', frame)
        if cv2.waitKey(1) & 0xFF == ord('q'): # 若按下Q键，则退出循环
            break
    else:
        break

cap.release()
out.release()
cv2.destroyAllWindows()


from django.urls import path
from .views import *

urlpatterns = [
    path('', home, name='home'),
    path('about/', about, name='about'),
    path('contact/', contact, name='contact'),
    path('register/', user_register, name='register'),
    path('login/', user_login, name='login'),
    path('dashboard/', dashboard, name='dashboard'),
    path('new_post/', new_post, name='new_post'),
    path('update_post/<int:pk>/', update_post, name='update_post'),
    path('delete_post/<int:pk>/', delete_post, name='delete_post'),
    path('logout/', user_logout, name='logout'),
]


"""Constants for the AMASTech integration."""
from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
import json, logging, asyncio, hashlib
import aiohttp
import async_timeout
from datetime import timedelta, datetime
from typing import Any
from homeassistant.components.sensor import SensorEntityDescription, SensorDeviceClass
from homeassistant.components.time import TimeEntityDescription
from homeassistant.components.number import NumberEntityDescription
from homeassistant.core import HomeAssistant
from homeassistant.components.binary_sensor import (
    BinarySensorDeviceClass,
    BinarySensorEntityDescription,
)
from homeassistant.const import TEMP_CELSIUS, PERCENTAGE
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
import os
from binascii import a2b_base64, b2a_base64, hexlify
from json import loads, dumps

_LOGGER = logging.getLogger(__name__)

DOMAIN = 'amas'
DEFAULT_NAME = 'AMAS'
DATA_KEY_API = 'api'
DATA_KEY_COORDINATOR = 'coordinator'
MIN_TIME_BETWEEN_UPDATES = timedelta(seconds=1)

def decrypt(payload, token):
     enc = a2b_base64(payload)
     IV = enc[:16]
     message = enc[16:]
     hmm = Cipher(algorithms.AES(token), modes.CBC(IV))
     decryptor = hmm.decryptor()
     data = decryptor.update(message) + decryptor.finalize()
     data = bytes((x for x in data if x >= 0x20 and x < 127))
     data = data.decode()
     return data

def encrypt(data, token):
    padder = padding.PKCS7(128).padder()
    data = padder.update(data) + padder.finalize()
    IV = os.urandom(16)
    cipher = Cipher(algorithms.AES(token), modes.CBC(IV))
    encryptor = cipher.encryptor()
    ct = encryptor.update(data) + encryptor.finalize()
    enc = IV + ct
    enc = b2a_base64(enc).decode().strip()
    return enc

def calculate_mac(enc):
	cbcmac = hashlib.sha256(enc.encode())
	cbcmac = hexlify(cbcmac.digest()).decode().strip()
	return cbcmac

# Encrypt and Calculate MAC
def encryptAndMac(message, token, mactoken):
	base64enc = encrypt(message, token)
	cbcmac = calculate_mac(base64enc)
	base64mac = encrypt(cbcmac.encode(), mactoken)
	return dumps({'base64enc': base64enc, 'base64mac': base64mac})

# Verify and Decrypt
def decryptAndVerify(enc_and_mac, token, mactoken):
	cbcmac = decrypt(enc_and_mac['base64mac'], mactoken)
	if calculate_mac(enc_and_mac['base64enc']) == cbcmac:
		return decrypt(enc_and_mac['base64enc'], token)
	else:
		return False

class AMASHub:
    """AMASHub class to check authentication and get device info.

    """

    def __init__(self, host: str, hass: HomeAssistant, session: aiohttp.ClientSession) -> None:
        """Initialize."""
        self.host = host
        self.session = session
        self.hass = hass
        self.loop = hass.loop
        self.api_key = ''
        self.mactoken = ''
        self.device_info = {}
        self.last_update = 0
        self.stream_task = None

    async def authenticate(self, api_key: str, mactoken: str) -> bool:
        """Test if we can decrypt responses."""
        url = 'http://' + self.host + '/control'
        try:
            api_key = a2b_base64(api_key)
            mactoken = a2b_base64(mactoken)
            body=loads(encryptAndMac(dumps({'state': {'desired': {}}}).encode(), api_key, mactoken))
            # r = requests.get(url, headers=headers)
            async with async_timeout.timeout(20):
                response = await self.session.post(url, json=body)
            if response.status == 200:
                payload = await response.json()
                _LOGGER.debug("Reponse content: %s", dumps(payload))
                try:
                    device_info = loads(decryptAndVerify(payload, api_key, mactoken))
                    device_info = device_info['state']['reported']
                except: raise ConfigEntryAuthFailed
                self.api_key = api_key
                self.mactoken = mactoken
                self.device_info = device_info
                self.last_update = datetime.now().strftime('%s')
                return True
            elif response.status == 500:
                 raise ConfigEntryNotReady
        except Exception as e:
            _LOGGER.warning("Failed to connect: %s", str(e))
            raise ConfigEntryNotReady

    async def check_connection(self) -> bool:
        """ Test if we can reconnect """
        url = 'http://' + self.host + '/control'
        try:
            body=loads(encryptAndMac(dumps({'state': {'desired': {}}}).encode(), self.api_key, self.mactoken))
            async with async_timeout.timeout(20):
                response = await self.session.post(url, json=body)
            if response.status == 200:
                payload = await response.json()
                _LOGGER.debug("Reponse content: %s", dumps(payload))
                try:
                    device_info = loads(decryptAndVerify(payload, self.api_key, self.mactoken))
                    device_info = device_info['state']['reported']
                except: raise ConfigEntryAuthFailed
                self.device_info = device_info
                self.last_update = datetime.now().strftime('%s')
                return True
            elif response.status == 500:
                 raise ConfigEntryNotReady
        except Exception as e:
            _LOGGER.warning("Failed to connect: %s", str(e))
            raise ConfigEntryNotReady

    async def control_device(self, state: dict[str, Any]) -> None:
        """Control device."""
        url = 'http://' + self.host + '/control'
        body = {'state': {'desired': state}}
        payload = loads(encryptAndMac(dumps(body).encode(), self.api_key, self.mactoken).encode())
        try:
            # r = requests.post(url, headers=headers, body=body)
            async with async_timeout.timeout(20):
                response = await self.session.post(url, json=payload)
                _LOGGER.debug("Response content: %s", str(response.content))
            if response.status == 200:
                device_info = await response.json()
                _LOGGER.debug("Reponse content: %s", str(device_info))
                try:
                    device_info = loads(decryptAndVerify(device_info, self.api_key,self.mactoken))
                except: raise ConfigEntryAuthFailed
                device_info = device_info['state']['reported']
                self.device_info = device_info
                self.last_update = datetime.now().strftime('%s')
                _LOGGER.debug("Device info: %s", str(response.content))
            else:
                _LOGGER.critical("Status code: "+str(response.status))
                raise ConfigEntryNotReady
        except Exception as e:
            _LOGGER.warning("Failed to connect: %s", str(e))
            raise ConfigEntryNotReady

    async def stream_info(self) -> None:
        url = 'http://' + self.host + '/metrics'
        try:
            async with self.session.ws_connect(url) as ws:
                async for msg in ws:
                    _LOGGER.debug('WSMsgType: ' + str(msg.type))
                    if msg.type == aiohttp.WSMsgType.ERROR:
                        await ws.close()
                    elif msg.type == aiohttp.WSMsgType.BINARY:
                        try:
                            device_info = loads(decryptAndVerify(loads(msg.data.decode()), self.api_key, self.mactoken))
                            device_info = device_info['state']['reported']
                            self.device_info = device_info
                            self.last_update = datetime.now().strftime('%s')
                        except: await ws.close()
                    elif msg.type == aiohttp.WSMsgType.TEXT:
                        try:
                            device_info = loads(decryptAndVerify(loads(msg.data), self.api_key, self.mactoken))
                            device_info = device_info['state']['reported']
                            self.device_info = device_info
                            self.last_update = datetime.now().strftime('%s')
                        except: await ws.close()
        except: _LOGGER.error('Streaming failed!')

@dataclass
class AMASNumberEntityDescription(NumberEntityDescription):
    """Describes AMASTech number entity."""

    icon: str = "mdi:leaf"

NUMBER_TYPES: tuple[AMASNumberEntityDescription, ...] = (
    AMASNumberEntityDescription(
        key="pump_interval",
        name="Pump Interval",
        icon="mdi:water-pump-off",
        native_max_value=86399,
        native_min_value=600,
    ),
    AMASNumberEntityDescription(
        key="pump_runtime",
        name="Pump Runtime",
        icon="mdi:water-pump",
        native_max_value=86399,
        native_min_value=60,
    ),
)

@dataclass
class AMASTimeEntityDescription(TimeEntityDescription):
    """Describes AMASTech time entity."""

    icon: str = "mdi:leaf"

TIME_TYPES: tuple[AMASTimeEntityDescription, ...] = (
    AMASTimeEntityDescription(
        key="light_on",
        name="Light On",
        icon="mdi:lightbulb-on",
    ),
    AMASTimeEntityDescription(
        key="light_off",
        name="Light Off",
        icon="mdi:lightbulb-off-outline",
    ),
)

@dataclass
class AMASSensorEntityDescription(SensorEntityDescription):
    """Describes AMASTech sensor entity."""

    icon: str = "mdi:leaf"

SENSOR_TYPES: tuple[AMASSensorEntityDescription, ...] = (
    AMASSensorEntityDescription(
        key="ambient_temperature",
        name="Ambient Temperature",
        native_unit_of_measurement=TEMP_CELSIUS,
        icon="mdi:thermometer",
        device_class=SensorDeviceClass.TEMPERATURE,
    ),
    AMASSensorEntityDescription(
        key="relative_humidity",
        name="Relative Humidity",
        native_unit_of_measurement=PERCENTAGE,
        icon="mdi:water-percent",
        device_class=SensorDeviceClass.HUMIDITY,
    ),
    AMASSensorEntityDescription(
        key="water_level",
        name="Water Level",
        native_unit_of_measurement=PERCENTAGE,
        entity_registry_enabled_default=False,
        icon="mdi:waves-arrow-up",
    ),
)

@dataclass
class RequiredAMASBinaryDescription:
    """Represent the required attributes of the AMASTech binary description."""

    state_value: Callable[[AMASHub], bool]

@dataclass
class AMASBinarySensorEntityDescription(
    BinarySensorEntityDescription, RequiredAMASBinaryDescription
):
    """Describes AMASTech binary sensor entity."""

    extra_value: Callable[[AMASHub], dict[str, Any] | None] = lambda api: None

BINARY_SENSOR_TYPES: tuple[AMASBinarySensorEntityDescription, ...] = (
    AMASBinarySensorEntityDescription(
        key="light_status",
        name="Light Status",
        entity_registry_enabled_default=True,
        device_class=BinarySensorDeviceClass.LIGHT,
        state_value=lambda api: bool(api.device_info['light']['status']),
    ),
    AMASBinarySensorEntityDescription(
        key="pump_status",
        name="Pump Status",
        entity_registry_enabled_default=True,
        device_class=BinarySensorDeviceClass.RUNNING,
        state_value=lambda api: bool(api.device_info['pump']['status']),
    ),
    AMASBinarySensorEntityDescription(
        key="water_level_alert",
        name="Water Level Alert",
        entity_registry_enabled_default=True,
        device_class=BinarySensorDeviceClass.BATTERY,
        state_value=lambda api: api.device_info['alerts']['water_level_alert'] == 'Low',
    ),
    AMASBinarySensorEntityDescription(
        key="amb_temp_alert",
        name="Ambient Temperature Alert",
        entity_registry_enabled_default=False,
        device_class=BinarySensorDeviceClass.PROBLEM,
        state_value=lambda api: api.device_info['alerts']['temp_alert'] == 'Low' or api.device_info['alerts']['temp_alert'] == 'High',
    ),
    AMASBinarySensorEntityDescription(
        key="rel_humidity_alert",
        name="Relative Humidity Alert",
        entity_registry_enabled_default=False,
        device_class=BinarySensorDeviceClass.PROBLEM,
        state_value=lambda api: api.device_info['alerts']['humidity_alert'] == 'Low' or api.device_info['alerts']['humidity_alert'] == 'High',
    ),
)


from beartype.typing import (
    Any,
    Callable,
    Union,
)
from jax import lax
from jax.experimental import host_callback
from jaxtyping import Float  # noqa: TCH002
from tqdm.auto import tqdm

from gpjax.typing import Array  # noqa: TCH001

def progress_bar(num_iters: int, log_rate: int) -> Callable:
    r"""Progress bar decorator for the body function of a `jax.lax.scan`.

    Example:
    ```python
        >>> import jax.numpy as jnp
        >>> import jax
        >>>
        >>> carry = jnp.array(0.0)
        >>> iteration_numbers = jnp.arange(100)

        >>> @progress_bar(num_iters=iteration_numbers.shape[0], log_rate=10)
        >>> def body_func(carry, x):
        ...    return carry, x
        >>>
        >>> carry, _ = jax.lax.scan(body_func, carry, iteration_numbers)
    ```

    Adapted from this [excellent blog post](https://www.jeremiecoullon.com/2021/01/29/jax_progress_bar/).

    Might be nice in future to directly create a general purpose `verbose scan`
    inplace of a for a jax.lax.scan, that takes the same arguments as a jax.lax.scan,
    but prints a progress bar.
    """
    tqdm_bars = {}
    remainder = num_iters % log_rate

    """Define a tqdm progress bar."""
    tqdm_bars[0] = tqdm(range(num_iters))
    tqdm_bars[0].set_description("Compiling...", refresh=True)

    def _update_tqdm(args: Any, transform: Any) -> None:
        r"""Update the tqdm progress bar with the latest objective value."""
        value, iter_num = args
        tqdm_bars[0].set_description("Running", refresh=False)
        tqdm_bars[0].update(iter_num)
        tqdm_bars[0].set_postfix({"Value": f"{value: .2f}"})

    def _close_tqdm(args: Any, transform: Any) -> None:
        r"""Close the tqdm progress bar."""
        tqdm_bars[0].close()

    def _callback(cond: bool, func: Callable, arg: Any) -> None:
        r"""Callback a function for a given argument if a condition is true."""
        dummy_result = 0

        def _do_callback(_) -> int:
            """Perform the callback."""
            return host_callback.id_tap(func, arg, result=dummy_result)

        def _not_callback(_) -> int:
            """Do nothing."""
            return dummy_result

        _ = lax.cond(cond, _do_callback, _not_callback, operand=None)

    def _update_progress_bar(value: Float[Array, "1"], iter_num: int) -> None:
        r"""Update the tqdm progress bar."""
        # Conditions for iteration number
        is_multiple: bool = (iter_num % log_rate == 0) & (
            iter_num != num_iters - remainder
        )
        is_remainder: bool = iter_num == num_iters - remainder
        is_last: bool = iter_num == num_iters - 1

        # Update progress bar, if multiple of log_rate
        _callback(is_multiple, _update_tqdm, (value, log_rate))

        # Update progress bar, if remainder
        _callback(is_remainder, _update_tqdm, (value, remainder))

        # Close progress bar, if last iteration
        _callback(is_last, _close_tqdm, None)

    def _progress_bar(body_fun: Callable) -> Callable:
        r"""Decorator that adds a progress bar to `body_fun` used in `jax.lax.scan`."""

        def wrapper_progress_bar(carry: Any, x: Union[tuple, int]) -> Any:
            # Get iteration number
            if type(x) is tuple:
                iter_num, *_ = x
            else:
                iter_num = x

            # Compute iteration step
            result = body_fun(carry, x)

            # Get value
            *_, value = result

            # Update progress bar
            _update_progress_bar(value, iter_num)

            return result

        return wrapper_progress_bar

    return _progress_bar

__all__ = [
    "progress_bar",
]


"""create locations table

Revision ID: 468946541a57
Revises: 26e3340ce799
Create Date: 2013-11-03 23:24:23.069426

"""
from alembic import op
import sqlalchemy as sa

revision = '468946541a57'
down_revision = '26e3340ce799'

def upgrade():
    op.create_table(
        'locations',
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('created', sa.DateTime(timezone=True)),
        sa.Column('city', sa.Unicode(75), nullable=False),
        sa.Column('country_code', sa.Unicode(3), nullable=False)
    )

def downgrade():
    op.drop_table('locations')


#----------------------------------------------------------------------------#
# Imports
#----------------------------------------------------------------------------#
import json
import dateutil.parser
import babel
from flask import Flask, render_template, request, Response, flash, redirect, url_for
from flask_moment import Moment
from flask_sqlalchemy import SQLAlchemy
import logging
from logging import Formatter, FileHandler
from flask_wtf import Form
from forms import *
from datetime import date
from utils import get_num_upcoming_shows_at_venue, get_num_upcoming_shows_for_artist
from models import *
#----------------------------------------------------------------------------#
# App Config.
#----------------------------------------------------------------------------#

# TODO: connect to a local postgresql database
# Create the database automatically if it doesn't exist
from sqlalchemy import create_engine
from sqlalchemy_utils import database_exists, create_database
from config import SQLALCHEMY_DATABASE_URI
engine = create_engine(SQLALCHEMY_DATABASE_URI)
if not database_exists(engine.url):
    create_database(engine.url)

app = Flask(__name__)
moment = Moment(app)
app.config.from_object('config')
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
db.init_app(app)

from flask_migrate import Migrate

migrate = Migrate(app, db)

# from add_initial_data import add_initial_data
# add_initial_data()

#----------------------------------------------------------------------------#
# Filters.
#----------------------------------------------------------------------------#

def format_datetime(value, format='medium'):
  date = dateutil.parser.parse(value)
  if format == 'full':
      format="EEEE MMMM, d, y 'at' h:mma"
  elif format == 'medium':
      format="EE MM, dd, y h:mma"
  return babel.dates.format_datetime(date, format)

app.jinja_env.filters['datetime'] = format_datetime

#----------------------------------------------------------------------------#
# Controllers.
#----------------------------------------------------------------------------#

@app.route('/')
def index():
  return render_template('pages/home.html')

#  Venues
#  ----------------------------------------------------------------

@app.route('/venues')
def venues():
  # TODO: replace with real venues data.
  #       num_shows should be aggregated based on number of upcoming shows per venue.
  state_cities = set()

  for state_city in db.session.query(Venue.state, Venue.city).all():
      state_cities.add(state_city)

  state_cities = sorted(list(state_cities))
  data = []
  for (state, city) in state_cities:
    data.append({})
    data[-1]["city"] = city
    data[-1]["state"] = state
    data[-1]["venues"] = []

    for venue_id, venue_name in db.session.query(Venue.id, Venue.name).filter(Venue.city==city):
      upcoming_shows = get_num_upcoming_shows_at_venue(db, venue_id, Show)
      venue_details = {}
      venue_details["id"] = venue_id
      venue_details["name"] = venue_name
      venue_details["num_upcoming_shows"] = upcoming_shows
      data[-1]["venues"].append(venue_details)

  return render_template('pages/venues.html', areas=data);

@app.route('/venues/search', methods=['POST'])
def search_venues():
  # TODO: implement search on venues with partial string search. Ensure it is case-insensitive.
  # seach for Hop should return "The Musical Hop".
  # search for "Music" should return "The Musical Hop" and "Park Square Live Music & Coffee"
  search_term = request.form.get('search_term', '')
  venues = db.session.query(Venue.id, Venue.name).all()
  response = {"count": 0, "data": []}

  for venue_id, venue_name in venues:
    # venue_id, venue_name = venue_ids[ix], venue_names[ix]
    if search_term.lower() in venue_name.lower() or venue_name.lower() in search_term.lower():
      response["count"] += 1
      upcoming_shows = get_num_upcoming_shows_at_venue(db, venue_id, Show)
      venue_info = {"id": venue_id, "name": venue_name, "num_upcoming_shows": upcoming_shows}
      response["data"].append(venue_info)
  return render_template('pages/search_venues.html', results=response, search_term=search_term)

@app.route('/venues/<int:venue_id>')
def show_venue(venue_id):
  # shows the venue page with the given venue_id
  # TODO: replace with real venue data from the venues table, using venue_id

  result = db.session.query(Venue).filter(Venue.id==venue_id)[0]
  data = result.__dict__
  data["genres"] =  data["genres"].split(",")

  data["past_shows"] = []
  data["upcoming_shows"] = []
  for show in db.session.query(Show).filter(Show.venue_id==venue_id).all():
      show_info = {}
      format="EE MM, dd, y h:mma"
      show_info["start_time"] =  babel.dates.format_datetime(show.start_time, format)
      artist_info = db.session.query(Artist.id, Artist.name, Artist.image_link).filter(Artist.id==show.artist_id)[0]
      show_info["artist_id"] = artist_info.id
      show_info["artist_name"] = artist_info.name
      show_info["artist_image_link"] = artist_info.image_link

      if show.start_time <= date.today():
          data["past_shows"].append(show_info)
      else:
          data["upcoming_shows"].append(show_info)

  data["past_shows_count"] = len(data["past_shows"])
  data["upcoming_shows_count"] = len(data["upcoming_shows"])

  return render_template('pages/show_venue.html', venue=data)

#  Create Venue
#  ----------------------------------------------------------------

@app.route('/venues/create', methods=['GET'])
def create_venue_form():
  form = VenueForm()
  return render_template('forms/new_venue.html', form=form)

@app.route('/venues/create', methods=['POST'])
def create_venue_submission():
  # TODO: insert form data as a new Venue record in the db, instead
  # TODO: modify data to be the data object returned from db insertion

  try:
    venue = Venue()
    venue.name = request.form["name"]
    venue.genres = request.form["genres"]
    venue.city = request.form["city"]
    venue.state = request.form["state"]
    venue.address = request.form["address"]
    venue.phone = request.form["phone"]
    venue.facebook_link = request.form["facebook_link"]
    db.session.add(venue)
    db.session.commit()
    flash('Venue ' + request.form['name'] + ' was successfully listed!')

  except:
    db.session.rollback()
    flash('An error occurred. Venue ' + request.form["name"] + ' could not be listed.')
    # TODO: on unsuccessful db insert, flash an error instead.
    # e.g., flash('An error occurred. Venue ' + data.name + ' could not be listed.')
    # see: http://flask.pocoo.org/docs/1.0/patterns/flashing/

  finally:
    db.session.close()

  return render_template('pages/home.html')

@app.route('/venues/<venue_id>', methods=['DELETE'])
def delete_venue(venue_id):
  # TODO: Complete this endpoint for taking a venue_id, and using
  # SQLAlchemy ORM to delete a record. Handle cases where the session commit could fail.
  try:
    venue = Venue.query.filter_by(venue_id==venue_id).one()
    db.session.delete(venue)
    db.session.commit()
    flash('Venue ' + venue_id + ' was successfully delete!')

  except:
    db.session.rollback()
    flash('An error occurred. Venue ' + venue_id + ' could not be deleted.')

  finally:
    db.session.close()

  # BONUS CHALLENGE: Implement a button to delete a Venue on a Venue Page, have it so that
  # clicking that button delete it from the db then redirect the user to the homepage
  return None

#  Artists
#  ----------------------------------------------------------------
@app.route('/artists')
def artists():
  # TODO: replace with real data returned from querying the database

  results = db.session.query(Artist.id, Artist.name).all()
  data = []
  for result in results:
    info = {"id": result.id, "name": result.name}
    data.append(info)
  return render_template('pages/artists.html', artists=data)

@app.route('/artists/search', methods=['POST'])
def search_artists():
  # TODO: implement search on artists with partial string search. Ensure it is case-insensitive.
  # seach for "A" should return "Guns N Petals", "Matt Quevado", and "The Wild Sax Band".
  # search for "band" should return "The Wild Sax Band".
  search_term = request.form.get('search_term', '')
  artists = db.session.query(Artist.id, Artist.name).all()
  response = {"count": 0, "data": []}
  for artist_id, artist_name in artists:
    if search_term.lower() in artist_name.lower() or artist_name.lower() in search_term.lower():
      response["count"] += 1
      upcoming_shows = get_num_upcoming_shows_for_artist(db, artist_id, Show)
      artist_info = {"id": artist_id, "name": artist_name, "num_upcoming_shows": upcoming_shows}
      response["data"].append(artist_info)
  return render_template('pages/search_artists.html', results=response, search_term=search_term)

@app.route('/artists/<int:artist_id>')
def show_artist(artist_id):
  # shows the venue page with the given venue_id
  # TODO: replace with real venue data from the venues table, using venue_id
  time_format="EE MM, dd, y h:mma"
  result = db.session.query(Artist).filter(Artist.id==artist_id)[0]
  data = result.__dict__
  data["genres"] =  data["genres"].split(",")

  data["past_shows"] = []
  data["upcoming_shows"] = []
  for show in db.session.query(Show).filter(Show.artist_id==artist_id).all():
      show_info = {}
      show_info["start_time"] =  babel.dates.format_datetime(show.start_time, time_format)
      venue_info = db.session.query(Venue.id, Venue.name, Venue.image_link).filter(Venue.id==show.venue_id)[0]
      show_info["venue_id"] = venue_info.id
      show_info["venue_name"] = venue_info.name
      show_info["venue_image_link"] = venue_info.image_link

      if show.start_time <= date.today():
          data["past_shows"].append(show_info)
      else:
          data["upcoming_shows"].append(show_info)

  data["past_shows_count"] = len(data["past_shows"])
  data["upcoming_shows_count"] = len(data["upcoming_shows"])
  return render_template('pages/show_artist.html', artist=data)

#  Update
#  ----------------------------------------------------------------
@app.route('/artists/<int:artist_id>/edit', methods=['GET'])
def edit_artist(artist_id):
  form = ArtistForm()
  artist = db.session.query(Artist).filter(Artist.id==artist_id).one().__dict__
  artist = {key: artist[key] for key in artist.keys() if key != "_sa_instance_state"}
  # TODO: populate form with fields from artist with ID <artist_id>
  return render_template('forms/edit_artist.html', form=form, artist=artist)

@app.route('/artists/<int:artist_id>/edit', methods=['POST'])
def edit_artist_submission(artist_id):
  # TODO: take values from the form submitted, and update existing
  # artist record with ID <artist_id> using the new attributes
  try:
    artist = Artist.query.filter(Artist.id==artist_id).first()
    new_info = request.form.__dict__

    for key in new_info.keys():
        setattr(artist, key, new_info[key])

    db.session.commit()
  except:
    db.session.rollback()
  finally:
    db.session.close()

  return redirect(url_for('show_artist', artist_id=artist_id))

@app.route('/venues/<int:venue_id>/edit', methods=['GET'])
def edit_venue(venue_id):
  form = VenueForm()
  venue = db.session.query(Venue).filter(Venue.id==venue_id).one().__dict__
  venue = {key: venue[key] for key in venue.keys() if key != "_sa_instance_state"}
  # TODO: populate form with values from venue with ID <venue_id>
  return render_template('forms/edit_venue.html', form=form, venue=venue)

@app.route('/venues/<int:venue_id>/edit', methods=['POST'])
def edit_venue_submission(venue_id):
  # TODO: take values from the form submitted, and update existing
  # venue record with ID <venue_id> using the new attributes
  try:
    venue = Venue.query.filter(Venue.id==venue_id).first()
    new_info = request.form.__dict__

    for key in new_info.keys():
        setattr(venue, key, new_info[key])
    db.session.commit()
  except:
    db.session.rollback()
  finally:
    db.session.close()

  return redirect(url_for('show_venue', venue_id=venue_id))

#  Create Artist
#  ----------------------------------------------------------------

@app.route('/artists/create', methods=['GET'])
def create_artist_form():
  form = ArtistForm()
  return render_template('forms/new_artist.html', form=form)

@app.route('/artists/create', methods=['POST'])
def create_artist_submission():
  # called upon submitting the new artist listing form
  # TODO: insert form data as a new Venue record in the db, instead
  # TODO: modify data to be the data object returned from db insertion
  try:
    artist = Artist()
    artist.name = request.form["name"]
    artist.genres = request.form["genres"]
    artist.city = request.form["city"]
    artist.state = request.form["state"]
    artist.phone = request.form["phone"]
    artist.facebook_link = request.form["facebook_link"]
    db.session.add(artist)
    db.session.commit()
    flash('Artist ' + request.form['name'] + ' was successfully listed!')

  except:
    db.session.rollback()
    flash('An error occurred. Artist ' + request.form["name"] + ' could not be listed.')

  finally:
    db.session.close()

  return render_template('pages/home.html')

#  Shows
#  ----------------------------------------------------------------

@app.route('/shows')
def shows():
  # displays list of shows at /shows
  # TODO: replace with real venues data.
  #       num_shows should be aggregated based on number of upcoming shows per venue.
  time_format="EE MM, dd, y h:mma"
  shows = db.session.query(Show).all()
  data = []
  for show in shows:
    result = {"venue_id": show.venue_id, "venue_name": show.venue.name, "artist_id": show.artist_id,
              "artist_image_link": show.artist.image_link, "start_time": babel.dates.format_datetime(show.start_time, time_format)}
    data.append(result)
  return render_template('pages/shows.html', shows=data)

@app.route('/shows/create')
def create_shows():
  # renders form. do not touch.
  form = ShowForm()
  return render_template('forms/new_show.html', form=form)

@app.route('/shows/create', methods=['POST'])
def create_show_submission():
  # called to create new shows in the db, upon submitting new show listing form
  # TODO: insert form data as a new Show record in the db, instead
  try:
    show = Show()
    show.artist_id = request.form["artist_id"]
    show.venue_id = request.form["venue_id"]
    show.start_time = request.form["start_time"]
    db.session.add(show)
    db.session.commit()
    flash('Show was successfully listed!')

  except:
    db.session.rollback()
    flash('An error occurred. Show not listed')

  finally:
    db.session.close()

  # on successful db insert, flash success

  # TODO: on unsuccessful db insert, flash an error instead.
  # e.g., flash('An error occurred. Show could not be listed.')
  # see: http://flask.pocoo.org/docs/1.0/patterns/flashing/
  return render_template('pages/home.html')

@app.errorhandler(404)
def not_found_error(error):
    return render_template('errors/404.html'), 404

@app.errorhandler(500)
def server_error(error):
    return render_template('errors/500.html'), 500

if not app.debug:
    file_handler = FileHandler('error.log')
    file_handler.setFormatter(
        Formatter('%(asctime)s %(levelname)s: %(message)s [in %(pathname)s:%(lineno)d]')
    )
    app.logger.setLevel(logging.INFO)
    file_handler.setLevel(logging.INFO)
    app.logger.addHandler(file_handler)
    app.logger.info('errors')

#----------------------------------------------------------------------------#
# Launch.
#----------------------------------------------------------------------------#

# Default port:
if __name__ == '__main__':
    app.run()

# Or specify port manually:
'''
if __name__ == '__main__':
    port = int(os.environ.get('PORT', 5000))
    app.run(host='0.0.0.0', port=port)
'''


import sys
import os
import json
import numpy as np
import pandas as pd

class Model:
    def __init__(self, inputPath):
        '''
        load dnn model with Keras
        '''
        self.modelPath = inputPath

        if os.path.exists("/".join([self.modelPath, "checkpoints"])):
            self.modelPath = "/".join([self.modelPath, "checkpoints"])
        # load config file
        configFile = self.modelPath+"/net_config.json"
        if not os.path.exists(configFile):
            sys.exit("could not find config file {}, cannot load DNN".format(configFile))

        with open(configFile) as f:
            self.config = f.read()
        self.config = json.loads(self.config)

        # load trained model
        checkpointPath = self.modelPath+"/trained_model.h5py"
        if not os.path.exists(checkpointPath):
            sys.exit("could not find trained model {}".format(checkpointPath))

        print("\nloading DNN model from {}\n".format(self.modelPath))
        try:
            from keras.models import load_model
            self.model = load_model(checkpointPath)
        except:
            from tensorflow.keras.models import load_model
            self.model = load_model(checkpointPath)

        self.model.summary()

    def setVariables(self):
        '''
        read list of variables used for DNN evaluation
        read norm table with mean and std.dev values for all of theses variables
        '''
        # load variable normalization
        normFile = self.modelPath+"/variable_norm.csv"
        if not os.path.exists(normFile):
            sys.exit("could not find file with variable norms {}".format(normFile))

        self.variable_norms = pd.read_csv(normFile, index_col=0)
        print("\nvariable norms:")
        print(self.variable_norms)
        self.variables = list(self.variable_norms.index.values)

    def findBest(self, entry):
        '''
        evaluate DNN for all entries given

        output: index, value of best combination
        '''
        matrix = entry[self.variables].values
        idx = 0
        for _, row in self.variable_norms.iterrows():
            matrix[:,idx]-=row["mu"]
            matrix[:,idx]/=row["std"]
            idx+=1
        predictions = self.model.predict(matrix)
        bestIndex = np.argmax(predictions[:,0])
        return bestIndex, predictions[bestIndex][0]

    def evaluate(self, entry):
        '''
        evaluate dnn with inputs
        inputs are either ndarray of events already cut to the dnn input variables
        or dataframe with correct input variable column names

        output: dnn predictions, max value of prediction vectors
        '''
        if not type(entry) == np.ndarray:
            matrix = entry[self.variables].values
        else:
            matrix = entry

        # norm variables
        idx=0
        for _, row in self.variable_norms.iterrows():
            matrix[:,idx]-=row["mu"]
            matrix[:,idx]/=row["std"]
            idx+=1

        # predict
        predictions = self.model.predict(matrix)

        # get maximum index
        maxIndex = np.argmax(predictions, axis = 1)
        return predictions, maxIndex


from django.http import HttpResponse
from django.shortcuts import render
from datetime import datetime
import requests
from django.shortcuts import render
from drf_yasg import openapi
from drf_yasg.utils import swagger_auto_schema
from rest_framework import status, views
from rest_framework.decorators import api_view, permission_classes
from rest_framework.permissions import AllowAny
from rest_framework.response import Response
from authentication.models import User

# imports from yako
from apirequests.apiurls import RadioQuote, ThirdPartyQoute, CheckVehicle, ThirdPartyPayment, ThirdPartyPolicy, ThirdPartyZinaraQoute, ZinaraQuote,LicensingPayment, LicensingPolicy
from config.models import InsuranceApiUrlConfig
from commission.models import AgentSale

# Third Party Only Quote
from config.models import InsuranceApiUrlConfig

# from apirequests.apiurls import ThirdPartyQoute, CheckVehicle, ThirdPartyPayment, ThirdPartyPolicy, RadioQuote,ZinaraQuote, LicensingPayment, LicensingPolicy
# from config.Endpoints import ThirdPartyQoute, ThirdPartyUpdate, ThirdPartyPolicy, LicensingQoute, LicensingPolicy, \
#     LicensingUpdate, CombinedPolicy, CombinedUpdate, CombinedQoute, ComprehensivePolicy, REVERSAL

#
# @permission_classes((AllowAny,))
# class TokenView(views.APIView):
#     def get(self, request):
#         response = requests.get(url=Yako + '/v2/icecash-token/')
#         if response.status_code == 200:
#             data = response.json()
#             token = data['token']
#             res = {"message": "%s" % token}
#             return Response(res, status=status.HTTP_200_OK)

class PolicyReversalView(views.APIView):

    def get(self, request):

        if request.method == 'GET':
            original_partner_ref = request.GET.get('original_partner_reference')
            try:
                quote = requests.get(
                    REVERSAL()[0]['api_endpoint'] + '?original_partner_reference=' + original_partner_ref )
                res = quote.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return HttpResponse(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

class ComprehensiveView(views.APIView):
    """
    Key for request_type (INT):
        1 =  Quote
        2 =  Update
        3 =  Retrieve Policy
    """

    def get(self, request):

        if request.method == 'GET':
            vrn = request.GET.get('vrn')
            durationMonths = request.GET.get('duration_months')
            vehicle_type = request.query_params.get('vrn_type', None)
            vrn_usage = request.query_params.get('vrn_usage', None)
            suminsured = request.query_params.get('suminsured', None)
            product = request.query_params.get('product', None)

            vrn_type_icecash = request.query_params.get('vrn_type_icecash', None)
            vrn_usage_icecash = request.query_params.get('vrn_usage_icecash', None)

            try:
                quote = requests.get('http://196.43.100.211:3005/comprehensive/quote/' + '?suminsured=' + suminsured + '&product=' + product +
                    '&vrn=' + vrn + '&vrn_type=' + vehicle_type + '&vrn_usage=' + vrn_usage + '&duration_months=' + durationMonths + '&vrn_usage_icecash=' + vrn_usage_icecash + '&vrn_type_icecash=' + vrn_type_icecash)
                res = quote.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return HttpResponse(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

class ThirdPartyView(views.APIView):
    """
    Key for request_type (INT):
        1 =  Quote
        2 =  Update
        3 =  Retrieve Policy
    """

    def get(self, request):
        print('ive arrived')
        request_type = request.GET.get('request_type')

        if request.method == 'GET' and request_type == '1':
            vrn = request.GET.get('vrn')
            durationMonths = request.GET.get('duration_months')
            vehicle_type = request.query_params.get('vrn_type', None)
            tax_class = request.query_params.get('tax_class', None)
            print(request_type, vrn, tax_class, durationMonths, vehicle_type, ThirdPartyQoute())
            print(ThirdPartyQoute() + '?vrn=' + vrn + '&duration_months=' + durationMonths +
                  '&vrn_type=' + vehicle_type + '&tax_class=' + tax_class)
            try:
                # quote = requests.get(
                #     ThirdPartyQoute() + '?vrn=' + vrn + '&duration_months=' + durationMonths +
                #     '&vrn_type=' + vehicle_type + '&tax_class=' + tax_class)

                quote = requests.get(ThirdPartyQoute() + '?vrn=' + vrn + '&duration_months=' + durationMonths + '&tax_class='+tax_class + '&vrn_type='+vehicle_type)
                res = quote.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return HttpResponse(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

        elif request.method == 'GET' and request_type == '2':
            paymentMethod = request.GET.get('payment_method')
            IDnumber = request.GET.get('client_idnumber')
            insuranceID = request.GET.get('insurance_id')
            payment_status = request.GET.get('payment_status')
            msisdn = request.GET.get('client_mobile')
            print(paymentMethod, IDnumber, insuranceID, payment_status, msisdn)
            try:
                payment_update = requests.get(ThirdPartyPayment() + '?payment_method=' + paymentMethod + '&client_idnumber=' + IDnumber + '&insurance_id=' + insuranceID +
                                              '&payment_status=' + payment_status + '&client_mobile=' + msisdn)
                res = payment_update.json()
                print(res, 'Json request')
                return HttpResponse(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return HttpResponse(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

        elif request.method == 'GET' and request_type == '3':
            insuranceID = request.GET.get('insuranceID')

            try:
                retrieve_policy = requests.get(ThirdPartyPolicy() + '?insuranceID=' + insuranceID)
                res = retrieve_policy.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return Response(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

class LicensingView(views.APIView):
    """
    Key for request_type (INT):
        1 =  Quote
        2 =  Update
        3 =  Retrieve Policy
    """

    def get(self, request):
        request_type = request.GET.get('request_type')

        if request.method == 'GET' and request_type == '1':
            print('im here')
            vrn = request.GET.get('vrn')
            IDnumber = request.GET.get('client_id_number')
            radio_usage = request.GET.get('radio_usage')
            LicFrequency = request.GET.get('frequency')
            radio_type = request.GET.get('type')
            # if(request.GET.get('radio_type') is not None ):
            #     radio_type = request.GET.get('radio_type')

            product_type = request.GET.get('type')
            vehicle_type = request.query_params.get('vrn_type', None)
            tax_class = request.query_params.get('tax_class', None)
            print('-------------------------------- DEBUG ---------------------------------')
            print(f"vrn: { vrn }, Lic Frequency: { LicFrequency }, ID_number: { IDnumber }, radio: { radio_type }, radio type: {radio_type}")
            print('-------------------------------- DEBUG ---------------------------------')
            try:
                # quote = requests.get(
                #     ThirdPartyZinaraQoute() + '?vrn=' + vrn + '&type=' + product_type + '&client_id_number='
                #     + IDnumber + '&radio_usage=' + radio_usage + '&frequency=' + LicFrequency
                #     + '&vrn_type=' + vehicle_type + '&tax_class=' + tax_class)

                quote = requests.get(RadioQuote() + '?vrn=' + vrn + '&frequency=' + LicFrequency + '&client_id_number='
                                     + IDnumber + "&radio_usage="+radio_usage + "&radio_type="+radio_type)
                res = quote.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return Response(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

        elif request.method == 'GET' and request_type == '2':
            paymentMethod = request.GET.get('payment_method')
            IDnumber = request.GET.get('client_idnumber')
            licenceID = request.GET.get('licence_id')
            payment_status = request.GET.get('payment_status')
            msisdn = request.GET.get('client_mobile')
            deliveryMethod = request.GET.get('delivery_method')
            try:
                # payment_update = requests.get(
                #     LicensingUpdate()[0]['api_endpoint'] + '?payment_method=' + paymentMethod +
                #     '&client_idnumber=' + IDnumber + '&licence_id=' + licenceID +
                #     '&payment_status' + payment_status + '&client_mobile='
                #     + msisdn + '&delivery_method=' + deliveryMethod)

                tran_update = requests.get(LicensingPayment() + '?payment_method=' + paymentMethod +
                                           '&client_idnumber=' + IDnumber + '&licence_id=' + licenceID +
                                           '&payment_status=' + payment_status + "&client_mobile=" + msisdn  + "&request_type=2" + "&delivery_method=" + deliveryMethod)
                res = tran_update.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return Response(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

        elif request.method == 'GET' and request_type == '3':
            licenceID = request.GET.get('licence_id')

            try:
                retrieve_policy = requests.get(LicensingPolicy() + '?licence_id=' + licenceID)
                res = retrieve_policy.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return HttpResponse(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

class CombinedView(views.APIView):
    """
    Key for request_type (INT):
        1 =  Quote
        2 =  Update
        3 =  Retrieve Policy
    """

    def get(self, request):
        request_type = request.GET.get('request_type')
        if request.method == 'GET' and request_type == '1':

            vrn = request.GET.get('vrn')
            idnumber = request.GET.get('id_number')
            LicFrequency = request.GET.get('frequency')
            durationMonths = request.GET.get('duration_months')
            radio_usage = request.query_params.get('radio_usage')
            vehicle_type = request.query_params.get('vrn_type', None)
            tax_class = request.query_params.get('tax_class', None)
            try:
                quote = requests.get(
                    CombinedQoute()[0]['api_endpoint'] + '?vrn=' + vrn + '&duration_months='
                    + durationMonths + '&id_number='
                    + idnumber + '&frequency=' + LicFrequency + '&radio_usage=' + radio_usage +
                    '&vrn_type=' + vehicle_type + '&tax_class=' + tax_class)
                res = quote.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return Response(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

        elif request.method == 'GET' and request_type == '2':

            paymentMethod = request.GET.get('payment_method')
            combinedid = request.GET.get('combined_id')
            payment_status = request.GET.get('payment_status')
            msisdn = request.GET.get('client_mobile')
            deliveryMethod = request.GET.get('delivery_method')
            try:
                payment_update = requests.get(CombinedUpdate()[0]['api_endpoint'] + '?payment_method=' + paymentMethod +
                                              '&combined_id=' + combinedid +
                                              '&payment_status' + payment_status + '&client_mobile='
                                              + msisdn + '&delivery_method=' + deliveryMethod)
                res = payment_update.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return Response(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

        elif request.method == 'GET' and request_type == '3':
            combinedID = request.GET.get('combined_id')

            try:
                retrieve_policy = requests.get(CombinedPolicy()[0]['api_endpoint'] + '?combined_id=' + combinedID)
                res = retrieve_policy.json()
                print(res, 'Json request')
                return Response(res, status=status.HTTP_200_OK)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Outrisk For Service"}
                return Response(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)

class ThirdPartyPolicyView(views.APIView):
    insuranceID_param_config = openapi.Parameter('insuranceID', in_=openapi.IN_QUERY,
                                                 type=openapi.TYPE_STRING)

    @swagger_auto_schema(manual_parameters=[insuranceID_param_config])
    def get(self, request):
        insuranceID = request.GET.get('insuranceID')
        if request.method == 'GET':
            try:
                quote = requests.get(ThirdPartyPolicy() + '?insuranceID=' + insuranceID)
                res = quote.json()
                print(res, 'Json request')
                if res['Response']['Status'] == 'Approved':
                    print('Qoutes Succesful lets create save the transaction now')
                    # Process Commission Sharing
                    owner = User.objects.get(id=request.user.id)
                    price_category = owner.institution.agent_category.price
                    commission = float(res['Response']['Amount']) * price_category

                    print("================================================= DEBUG ====================================================")
                    print(f"owner: {owner}, price_category: {price_category}, commision {commission}")
                    print("================================================= DEBUG ====================================================")

                    AgentSale.objects.create(
                        agent=request.user,
                        agent_category=owner.institution.agent_category.category,
                        agent_institution=owner.institution.name,
                        agent_pricing=owner.institution.agent_category.price,
                        product_name='Third Party Insurance (RTA)',
                        vrn=res['Response']['VRN'],
                        transaction_id=insuranceID,
                        policy_number=res['Response']['PolicyNo'],
                        # receipt_number = res['Response']['ReceiptID'],
                        transaction_amount=float(res['Response']['Amount']),
                        transaction_commission=round(commission, 2),
                        transaction_status=res['Response']['Status'],
                        commission_month=datetime.now().strftime("%m-%Y")
                        # transaction_date = models.DateTimeField(auto_now=True),
                    )
                    return Response(res, status=status.HTTP_200_OK)
                else:
                    message = {"message": res}
                    return Response(message, status=status.HTTP_202_ACCEPTED)

            except ConnectionError as e:
                res = {"message": "Service Unavailable, Contact Xarani For Service"}
                return Response(res, status=status.HTTP_503_SERVICE_UNAVAILABLE)