# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""evaluate"""
import numpy as np

def eval_market1501(distmat, q_vids, g_vids, q_camids, g_camids, max_rank):
    """
    Evaluation with Market1501 metrics
    Key: for each query identity, its gallery images from the same camera view are discarded.
    """
    num_q, num_g = distmat.shape
    if num_g < max_rank:
        max_rank = num_g
        print("Note: number of gallery samples is quite small, got {}".format(num_g))
    indices = np.argsort(distmat, axis=1)
    matches = (g_vids[indices] == q_vids[:, np.newaxis]).astype(np.int32)

    # compute cmc curve for each query
    all_cmc = []
    all_AP = []
    num_valid_q = 0. # number of valid query
    for q_idx in range(num_q):
        # get query vid and camid
        q_vid = q_vids[q_idx]
        q_camid = q_camids[q_idx]

        # remove gallery samples that have the same vid and camid with query
        order = indices[q_idx]
        remove = (g_vids[order] == q_vid) & (g_camids[order] == q_camid)
        keep = np.invert(remove)

        # compute cmc curve
        orig_cmc = matches[q_idx][keep] # binary vector, positions with value 1 are correct matches
        if not np.any(orig_cmc):
            # this condition is true when query identity does not appear in gallery
            continue

        cmc = orig_cmc.cumsum()
        cmc[cmc > 1] = 1

        all_cmc.append(cmc[:max_rank])
        num_valid_q += 1.

        # compute average precision
        num_rel = orig_cmc.sum()
        tmp_cmc = orig_cmc.cumsum()
        tmp_cmc = [x / (i+1.) for i, x in enumerate(tmp_cmc)]
        tmp_cmc = np.asarray(tmp_cmc) * orig_cmc
        AP = tmp_cmc.sum() / num_rel
        all_AP.append(AP)

    assert num_valid_q > 0, "Error: all query identities do not appear in gallery"

    all_cmc = np.asarray(all_cmc).astype(np.float32)
    all_cmc = all_cmc.sum(0) / num_valid_q
    mAP = np.mean(all_AP)

    return all_cmc, mAP

def evaluate(distmat, q_vids, g_vids, q_camids, g_camids, max_rank=50, use_cython=True):
    return eval_market1501(distmat, q_vids, g_vids, q_camids, g_camids, max_rank)

def test(model, keyptaware, multitask, queryloader, galleryloader,
         vcolor2label, vtype2label, ranks=range(1, 51), return_distmat=False):
    """function eval"""
    model.set_train(False)

    qf = []
    q_vids = []
    q_camids = []
    q_vcolors = []
    q_vtypes = []
    pred_q_vcolors = []
    pred_q_vtypes = []
    for _, data in enumerate(queryloader.create_dict_iterator()):
        imgs = data["img"]
        vids = data["vid"]
        camids = data["camid"]
        vcolors = data["vcolor"]
        vtypes = data["vtype"]
        vkeypts = data["vkeypt"]

        _, output_vcolors, output_vtypes, features = model(imgs, vkeypts)

        qf.append(features.asnumpy())
        q_vids.extend(vids.asnumpy())
        q_camids.extend(camids.asnumpy())
        q_vcolors.extend(vcolors.asnumpy())
        q_vtypes.extend(vtypes.asnumpy())
        pred_q_vcolors.extend(output_vcolors.asnumpy())
        pred_q_vtypes.extend(output_vtypes.asnumpy())

    qf = np.concatenate(qf, axis=0)
    q_vids = np.asarray(q_vids)
    q_camids = np.asarray(q_camids)
    q_vcolors = np.asarray(q_vcolors)
    q_vtypes = np.asarray(q_vtypes)
    pred_q_vcolors = np.asarray(pred_q_vcolors)
    pred_q_vtypes = np.asarray(pred_q_vtypes)

    print("Extracted features for query set, obtained {}-by-{} matrix".format(qf.shape[0], qf.shape[1]))

    gf = []
    g_vids = []
    g_camids = []
    g_vcolors = []
    g_vtypes = []
    pred_g_vcolors = []
    pred_g_vtypes = []
    for _, data in enumerate(galleryloader.create_dict_iterator()):
        imgs = data["img"]
        vids = data["vid"]
        camids = data["camid"]
        vcolors = data["vcolor"]
        vtypes = data["vtype"]
        vkeypts = data["vkeypt"]

        _, output_vcolors, output_vtypes, features = model(imgs, vkeypts)

        gf.append(features.asnumpy())
        g_vids.extend(vids.asnumpy())
        g_camids.extend(camids.asnumpy())
        g_vcolors.extend(vcolors.asnumpy())
        g_vtypes.extend(vtypes.asnumpy())
        pred_g_vcolors.extend(output_vcolors.asnumpy())
        pred_g_vtypes.extend(output_vtypes.asnumpy())

    gf = np.concatenate(gf, axis=0)
    g_vids = np.asarray(g_vids)
    g_camids = np.asarray(g_camids)
    g_vcolors = np.asarray(g_vcolors)
    g_vtypes = np.asarray(g_vtypes)
    pred_g_vcolors = np.asarray(pred_g_vcolors)
    pred_g_vtypes = np.asarray(pred_g_vtypes)

    print("Extracted features for gallery set, obtained {}-by-{} matrix".format(gf.shape[0], gf.shape[1]))

    m, n = qf.shape[0], gf.shape[0]
    distmat = np.broadcast_to(np.power(qf, 2).sum(axis=1, keepdims=True), (m, n)) + \
              np.broadcast_to(np.power(gf, 2).sum(axis=1, keepdims=True), (n, m)).T

    distmat = distmat * 1 + (-2) * (np.matmul(qf, gf.T))

    print("Computing CMC and mAP")
    cmc, mAP = evaluate(distmat, q_vids, g_vids, q_camids, g_camids)

    print("Results ----------")
    print("mAP: {:.2%}".format(mAP))
    print("CMC curve")
    for r in ranks:
        print("Rank-{:<3}: {:.2%}".format(r, cmc[r-1]))
    print("------------------")

    print("Compute attribute classification accuracy")

    for q in range(q_vcolors.size):
        q_vcolors[q] = vcolor2label[q_vcolors[q]]
    for g in range(g_vcolors.size):
        g_vcolors[g] = vcolor2label[g_vcolors[g]]
    q_vcolor_errors = np.argmax(pred_q_vcolors, axis=1) - q_vcolors
    g_vcolor_errors = np.argmax(pred_g_vcolors, axis=1) - g_vcolors
    vcolor_error_num = np.count_nonzero(q_vcolor_errors) + np.count_nonzero(g_vcolor_errors)
    vcolor_accuracy = 1.0 - (float(vcolor_error_num) / float(distmat.shape[0] + distmat.shape[1]))
    print("Color classification accuracy: {:.2%}".format(vcolor_accuracy))

    for q in range(q_vtypes.size):
        q_vtypes[q] = vcolor2label[q_vtypes[q]]
    for g in range(g_vtypes.size):
        g_vtypes[g] = vcolor2label[g_vtypes[g]]
    q_vtype_errors = np.argmax(pred_q_vtypes, axis=1) - q_vtypes
    g_vtype_errors = np.argmax(pred_g_vtypes, axis=1) - g_vtypes
    vtype_error_num = np.count_nonzero(q_vtype_errors) + np.count_nonzero(g_vtype_errors)
    vtype_accuracy = 1.0 - (float(vtype_error_num) / float(distmat.shape[0] + distmat.shape[1]))
    print("Type classification accuracy: {:.2%}".format(vtype_accuracy))

    if return_distmat:
        return distmat
    return cmc[0]


# coding: utf8
import numpy as np
import argparse
import math
from time import clock, sleep
from utils.viewerClient import viewerClient
from solo8 import Solo8

def example_script(name_interface):
    viewer = viewerClient()
    device = Solo8(name_interface,dt=0.001)
    nb_motors = device.nb_motors

    q_viewer = np.array((7 + nb_motors) * [0.,])

    device.Init(calibrateEncoders=False)
    #CONTROL LOOP ***************************************************
    while ((not device.hardware.IsTimeout()) and (clock() < 200)):
        device.UpdateMeasurment()
        device.SetDesiredJointTorque([0]*nb_motors)
        device.SendCommand(WaitEndOfCycle=True)
        if ((device.cpt % 100) == 0):
            device.Print()

        q_viewer[3:7] = device.baseOrientation  # IMU Attitude
        q_viewer[7:] = device.q_mes  # Encoders
        viewer.display(q_viewer)
    #****************************************************************

    # Whatever happened we send 0 torques to the motors.
    device.SetDesiredJointTorque([0]*nb_motors)
    device.SendCommand(WaitEndOfCycle=True)

    if device.hardware.IsTimeout():
        print("Masterboard timeout detected.")
        print("Either the masterboard has been shut down or there has been a connection issue with the cable/wifi.")
    device.hardware.Stop()  # Shut down the interface between the computer and the master board

def main():
    parser = argparse.ArgumentParser(description='Example masterboard use in python.')
    parser.add_argument('-i',
                        '--interface',
                        required=True,
                        help='Name of the interface (use ifconfig in a terminal), for instance "enp1s0"')

    example_script(parser.parse_args().interface)

if __name__ == "__main__":
    main()


#!/usr/bin/python3

import cobra as cb
import pandas as pd
import re
import sys
import getopt
import os.path
import copy
import csv
import math
import cobra.flux_analysis.variability
import subprocess
import shutil, errno
import statistics
from cobra import Reaction
import cometspy as c
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams.update({'figure.max_open_warning': 0})
################################################################

def model_modifications(glu1, ac1, o21, glu2, ac2, o22):
    # Single GEMs parameter modifications
  # ===================================

    # 1.1.- [COBRApy] Establish modifications in model 1
    model1=c.model('/home/ana/FLYCOP/Scripts/iJO1366.cmd')
    model1.change_bounds('GLCtex_copy2', 0, 0)
    model1.change_bounds('PFK_3', 0, 0)
    model1.change_bounds('EX_o2_e',-1000,1000)
    model1.change_bounds('EX_glc__D_e',-1000,1000)
    model1.change_bounds('GLCtex_copy1', 0, -glu1)
    model1.change_bounds('O2tex', 0,-o21)
    if ac1<=0:
        model1.change_bounds('ACtex',-1000,-ac1)
        model1.change_bounds('EX_ac_e', ac1, 1000)
    else:
        model1.change_bounds('ACtex', -ac1,1000)
        model1.change_bounds('EX_ac_e',-1000,ac1)
    model1.id='ecoli_1_tmp'
    model1.write_comets_model()
    del(model1)

    # 1.2.- [COBRApy] Establish modifications in model 2
    model2=c.model("/home/ana/FLYCOP/Scripts/iJO1366.cmd")
    model2.change_bounds('GLCtex_copy2', 0, 0)
    model2.change_bounds('PFK_3', 0, 0)
    model2.change_bounds('EX_o2_e',-1000,1000)
    model2.change_bounds('EX_glc__D_e',-1000,1000)
    model2.change_bounds('GLCtex_copy1',0,-glu2)
    model2.change_bounds('O2tex',0,-o22)
    if(ac2<=0):
        model2.change_bounds('ACtex',-1000,-ac2)
        model2.change_bounds('EX_ac_e',ac2,1000)

    else:
        model2.change_bounds('ACtex',-ac2,1000)
        model2.change_bounds('EX_ac_e',-1000,ac2)
    model2.id='ecoli_2_tmp'
    model2.write_comets_model()
    del(model2)
    return (print('Modifications of the model done!'))

def initialize_layout(file):
    #The layout is a file with the stablished format
    layout=c.layout(file)
    return layout

def initialize_params(package, globall):
    """Function to initialize the comets' params class
    it can be initialize by submitting two files, one for the package parameters
    and one for the global ones.
    If you don't submit a file, the params class will be initialize with the values stated below
    which have been tested in this simulation"""

    if package or globall is not None:
        params = c.params(global_params = globall, package_params= package)
    else:
        params = c.params()
        params.all_params['maxCycles']=240
        params.all_params['timeStep']=0.1
        params.all_params['spaceWidth']=0.05
        params.all_params['allowCellOverlap']= True
        params.all_params['deathRate']= 0.01
        params.all_params['numRunThreads']= 4
        params.all_params['maxSpaceBiomass']= 10
        params.all_params['defaultVmax']=20
        params.all_params['showCycleTime']=True
        params.all_params['writeTotalBiomassLog']=True
        params.all_params['writeMediaLog']=True
        params.all_params['writeFluxLog']=True
        params.all_params['useLogNameTimeStamp']=False
        params.all_params['FluxLogRate']=1
        params.all_params['MediaLogRate']=1

    return params

def make_graph(comets, met, met22):

    '''This function creates a figure and saves it to pdf format.
    It also creates the file biomass_vs_met.txt which contais the quantity
    of each strain and metabolite and has the following columns:
    time(h), strain1 ... strainX, met1 ... metX.'''

    df =comets.media
    df=df.drop(columns=['x', 'y'])
    met2 =df[df['metabolite'] == str(met22)]
    met1 =df[df['metabolite'] == str(met)]
    met1 =met1.reset_index()
    met1=met1.drop(columns=['index', 'metabolite'])

    i = 45
    while i < 242:
        met1.loc[i] = [i+2, 0]
        i +=1

    met1=met1.drop(columns='cycle')

    met2=met2.reset_index()
    met2=met2.drop(columns=['metabolite', 'index'])

    j=46

    while j < 240:
        met2.loc[j] = [j+2, 0]
        j=j+1
    met2.loc[-1] = [1, 0]
    met2 =met2.sort_values(by=['cycle'])
    met2=met2.drop(columns='cycle')
    met2=met2.reset_index()
    met2=met2.drop(columns='index')
    df2=comets.total_biomass

    df2[str(met)]=met1
    df2[str(met22)]=met2

    df2.columns=['time(h)','Ecoli1', 'Ecoli2', str(met), str(met22)]
    np.savetxt(r'biomass_vs_glc_D_e_ac_e_template.txt', df2.values, fmt='%s',delimiter='\t')

    #reading the biomass file as a dataframe and adding it the metabolite concentrations
    plt.ioff()
    fig, ax = plt.subplots()
    ax.axis(ymax=1.6)
    ax.set_xlabel('time (h)')
    ax.set_ylabel('biomass (g/L)')
    l1,=ax.plot(df2['time(h)']*0.1, df2['Ecoli1'], label='Ecoli1', color='k')
    l2,=ax.plot(df2['time(h)']*0.1, df2['Ecoli2'], linestyle='--', label='Ecoli2', color='k')
    ax2 = ax.twinx()
    ax2.set_ylabel('metabolite conc (mM)')
    ax2.axis(ymin=0, ymax=10)
    l3,=ax2.plot(df2['time(h)']*0.1, df2[str(met)], label=str(met), color='b')
    l4,=ax2.plot(df2['time(h)']*0.1, df2[str(met22)], label=str(met22), color='r')
    fig.tight_layout()
    plt.legend([l1, l2, l3, l4],["Ecoli1", "Ecoli2", str(met), str(met22)])

    #Saving the graph as a pdf
    plt.savefig('biomass_vs_glc_D_e_ac_e_template_plot1.pdf')
    return df2
    #df2.drop(df2.index, inplace=True)

def end_simulation_cycle(ac1, ac2, df2, met):
    #function that stablishes the endCyle after a certain stop condition

    iniBiomass=float(df2.at[0,'Ecoli1'])+float(df2.at[0, 'Ecoli2'])
    totGlc=float(df2.at[0,'glc__D_e'])
    endGlcCycle=0
    for row in df2.iterrows():
        endCycle=int(row[1][0])
        glcConc=float(row[1][3])
        acConc=float(row[1][4])
        if endGlcCycle==0 and glcConc==0.0:
            endGlcCycle=endCycle
        if glcConc==0.0 and acConc==0.0:
            break
        if glcConc==0.0 and ac1>=0 and ac2>=0:
            break

    return iniBiomass, endCycle, totGlc

def biomass_yield(df2,endCycle):
    # Function that compute final biomass as the maximum biomass of each strain

    finalBiomass1=0
    finalBiomass2=0
    count=0
    for row in df2.iterrows():
        if float(row[1][1])>finalBiomass1:
            finalBiomass1=float(row[1][1])
        if float(row[1][2])>finalBiomass2:
            finalBiomass2=float(row[1][2])
        if count > endCycle:
            break
        count+=1
    finalBiomass=finalBiomass1+finalBiomass2
    return finalBiomass

def compute_fitness(biomassYieldNew, MaximumYield, finalBiomass, iniBiomass):
    #Function that computes the fitness according to the fitness function 'yield'
    fitness=(biomassYieldNew/MaximumYield) # Normalizing Yield
    if(float(finalBiomass-iniBiomass) > 1.03):  # Given that with both strains with WT, total biomass=1.028
        fitness=0
    return fitness

def uptakeMetabolite(comets,layout):
    #Function that calculates the uptake of the metabolite 'ac_e'
    #Changing the metabolite needs changing numRxnMet; the position of the reaction 'Ex_metName' in the model
    flux = comets.fluxes
    numRxnMet=88
    uptakeMet=[]
    i=0
    while i < len(layout.models)+1:
        uptakeMet.append(float(flux.iat[i+2,numRxnMet+3]))
        i+=1
    return uptakeMet

def ecoliLongTermFLYCOP_oneConf(glu1,ac1,o21,glu2,ac2,o22, fitFunc='MaxYield_MinTime', dirPlot='',repeat=10, params_package=None, params_global=None, layout_file='/home/ana/FLYCOP/Scripts/Nuevo_layout.txt',met='glc__D_e', met2='ac_e'):
    """Primarily function of the simulation"""

    print("Fitness function:"+fitFunc)

    if not(os.path.exists('/home/ana/FLYCOP/Scripts/ecoli_1_tmp.cmd')):
        model_modifications(glu1, ac1, o21, glu2, ac2, o22)

    layout=initialize_layout(layout_file)

    params =initialize_params(params_package, params_global)

    #establish comets class
    comets = c.comets(layout, params)

    """You may need to adjust the comets.set_classpath object if your programs aren't in the default comets location"""

    if not(os.path.exists('IndividualRunsResults')):
        os.makedirs('IndividualRunsResults')
    totfitness=0
    sumTotBiomass=0
    sumTotYield=0
    fitnessList=[]

    # To repeat X times, due to random behaviour in COMETS:
    for i in range(repeat):

        comets.run(delete_files=True)

        #obtaining the results and writing them to csv
        comets.total_biomass.to_csv('Total_biomass_log_template.txt',sep='\t',index=False)
        comets.fluxes.to_csv('Flux_log_template.txt',sep='\t', index=False)
        comets.media.to_csv('Media_log_template.txt',sep='\t', index=False)

        #graphic part
        df2 =make_graph(comets, met, met2)

        # 7.- Compute fitness (measure to optimize):
        print('computing fitness...')

        # 7.1.- Determine endCycle: when glucose and acetate are exhausted
        iniBiomass, endCycle, totGlc = end_simulation_cycle(ac1, ac2, df2, met)

        # 7.2.- Compute first element fitness: maximize biomass yield
        finalBiomass = biomass_yield(df2,endCycle)

        biomassYieldNew=float((finalBiomass-iniBiomass)/(totGlc*0.1801559)) # molecular weigth of met1 per mmol

        # For normalizing yield
        MaximumYield=0.6

        # 7.3.- Compute second element fitnes: minimize time
        fitTime=1-(float(endCycle)/float(240))

        # 7.4.- Compute joint fitness, as a 50% each element.
        fitness= compute_fitness(biomassYieldNew, MaximumYield, finalBiomass, iniBiomass)

        # 7.5.- Calculate the uptake of ac_e
        uptakeMet = uptakeMetabolite(comets,layout)

        print(" Total biomass: "+str(round(finalBiomass,6))+" in cycle "+str(endCycle)+". Biomass yield="+str(round(biomassYieldNew,6)))
        totfitness=totfitness+fitness
        fitnessList.append(fitness)
        sumTotBiomass=sumTotBiomass+finalBiomass
        sumTotYield=sumTotYield+biomassYieldNew

        # Copy individual solution
        file='IndividualRunsResults/'+'biomass_vs_glc_D_ac_run'+str(i)+'_'+str(fitness)+'_'+str(endCycle)+'.pdf'
        shutil.move('biomass_vs_glc_D_e_ac_e_template_plot1.pdf',file)
        if(dirPlot != ''):
            file2=dirPlot+'biomass_vs_glc_D_ac_'+str(glu1)+'_'+str(ac1)+'_'+str(o21)+'_'+str(glu2)+'_'+str(ac2)+'_'+str(o22)+'_'+str(round(uptakeMet[0],1))+'_'+str(round(uptakeMet[1],1))+'_run'+str(i)+'_'+str(fitness)+'_'+str(endCycle)+'.pdf'
            shutil.copy(file,file2)

        file='IndividualRunsResults/'+'total_biomass_log_run'+str(i)+'.txt'
        shutil.move('Total_biomass_log_template.txt',file)
        file='IndividualRunsResults/'+'media_log_run'+str(i)+'.txt'
        shutil.move('Media_log_template.txt',file)
        file='IndividualRunsResults/'+'flux_log_run'+str(i)+'.txt'
        shutil.move('Flux_log_template.txt',file)

    avgfitness=totfitness/repeat
    sdfitness=statistics.stdev(fitnessList)
    avgBiomass=sumTotBiomass/repeat
    avgYield=sumTotYield/repeat
    print("Fitness_function\tconfiguration\tfitness\tsd\tavg.Biomass\tavg.Yield\tendCycle")
    print(fitFunc+"\t"+str(glu1)+','+str(ac1)+','+str(o21)+','+str(glu2)+','+str(ac2)+','+str(o22)+','+str(round(uptakeMet[0],1))+','+str(round(uptakeMet[1],1))+"\t"+str(round(avgfitness,6))+"\t"+str(sdfitness)+"\t"+str(round(avgBiomass,6))+"\t"+str(round(avgYield,6))+"\t"+str(endCycle))
    with open(dirPlot+"configurationsResults"+fitFunc+".txt", "a") as myfile:
        myfile.write("Fitness_function\tconfiguration\tfitness\tsd\tavg.Biomass\tavg.Yield\tendCycle\n")
        myfile.write(fitFunc+"\t"+str(glu1)+','+str(ac1)+','+str(o21)+','+str(glu2)+','+str(ac2)+','+str(o22)+','+str(round(uptakeMet[0],1))+','+str(round(uptakeMet[1],1))+"\t"+str(round(avgfitness,6))+"\t"+str(sdfitness)+"\t"+str(round(avgBiomass,6))+"\t"+str(round(avgYield,6))+"\t"+str(endCycle)+"\n")
    print("Avg.fitness(sd):\t"+str(avgfitness)+"\t"+str(sdfitness)+"\n")
    if(sdfitness>0.1):
        avgfitness=0.0
    try:
        os.remove('/home/ana/FLYCOP/Scripts/ecoli_1_tmp.cmd')
    except FileNotFoundError:
        pass
    try:
        os.remove('/home/ana/FLYCOP/Scripts/ecoli_2_tmp.cmd')
    except FileNotFoundError:
        pass
    return avgfitness,sdfitness


import argparse
import os
import random

import cv2

import transforms
from spacenet_dataset import SpacenetLocDataset
from tools.metrics import calculate_metrics

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"

import zoo

from albumentations import Compose, PadIfNeeded, RandomBrightnessContrast, \
    RandomGamma, RandomRotate90, IAAAdditiveGaussianNoise, MedianBlur, DualTransform, HorizontalFlip, Resize

import losses

from apex.parallel import DistributedDataParallel, convert_syncbn_model
from tensorboardX import SummaryWriter

from tools.config import load_config
from tools.utils import create_optimizer, AverageMeter

from apex import amp

from losses import dice_round

import numpy as np
import torch
from torch.backends import cudnn
from torch.nn import DataParallel
from torch.utils.data import DataLoader
from tqdm import tqdm
import torch.distributed as dist

torch.backends.cudnn.benchmark = True

def create_val_transforms(conf):
    return Compose([
        PadIfNeeded(min_height=928, min_width=928)
    ])

def main():
    parser = argparse.ArgumentParser("PyTorch Spacenet6 Pipeline")
    arg = parser.add_argument
    arg('--config', metavar='CONFIG_FILE', help='path to configuration file')
    arg('--workers', type=int, default=8, help='number of cpu threads to use')
    arg('--gpu', type=str, default='0', help='List of GPUs for parallel training, e.g. 0,1,2,3')
    arg('--output-dir', type=str, default='weights/')
    arg('--resume', type=str, default='')
    arg('--fold', type=int, default=0)
    arg('--prefix', type=str, default='localization_')
    arg('--data-dir', type=str, default="/mnt/sota/datasets/spacenet/train/AOI_11_Rotterdam/")
    arg('--folds-csv', type=str, default='folds.csv')
    arg('--logdir', type=str, default='logs')
    arg('--zero-score', action='store_true', default=False)
    arg('--from-zero', action='store_true', default=False)
    arg('--distributed', action='store_true', default=False)
    arg("--local_rank", default=0, type=int)
    arg("--opt-level", default='O1', type=str)
    arg("--predictions", default="/wdata/oof_preds", type=str)
    arg("--test_every", type=int, default=1)
    arg("--visualizer-path", type=str, default="visualizer/visualizer.jar")

    args = parser.parse_args()

    if args.distributed:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(backend='nccl', init_method='env://')
    else:
        os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
        os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

    cudnn.benchmark = True

    conf = load_config(args.config)
    model = zoo.__dict__[conf['network']](seg_classes=conf['num_classes'], backbone_arch=conf['encoder'])

    model = model.cuda()
    if args.distributed:
        model = convert_syncbn_model(model)
    mask_loss_function = losses.__dict__[conf["mask_loss"]["type"]](**conf["mask_loss"]["params"]).cuda()
    loss_functions = {"mask_loss": mask_loss_function}
    optimizer, scheduler = create_optimizer(conf['optimizer'], model)

    dice_best = 0
    f_score_best = 0
    start_epoch = 0
    batch_size = conf['optimizer']['batch_size']

    data_train = SpacenetLocDataset(mode="train",
                                    fold=args.fold,
                                    data_path=args.data_dir,
                                    folds_csv=args.folds_csv,
                                    transforms=transforms.__dict__[conf["transforms"]](conf["input"]),
                                    multiplier=conf["data_multiplier"])
    data_val = SpacenetLocDataset(mode="val",
                                  fold=args.fold,
                                  data_path=args.data_dir,
                                  folds_csv=args.folds_csv,
                                  transforms=create_val_transforms(conf['input'])
                                  )
    train_sampler = None
    if args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(data_train)

    train_data_loader = DataLoader(data_train, batch_size=batch_size, num_workers=args.workers,
                                   shuffle=train_sampler is None, sampler=train_sampler, pin_memory=False,
                                   drop_last=True)
    val_batch_size = 1
    val_data_loader = DataLoader(data_val, batch_size=val_batch_size, num_workers=args.workers, shuffle=False,
                                 pin_memory=False)
    snapshot_name = "{}{}_{}_{}".format(args.prefix, conf['network'], conf['encoder'], args.fold)

    os.makedirs(args.logdir, exist_ok=True)
    summary_writer = SummaryWriter(args.logdir + '/' + snapshot_name)
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume, map_location='cpu')
            state_dict = checkpoint['state_dict']
            if conf['optimizer'].get('zero_decoder', False):
                for key in state_dict.copy().keys():
                    if key.startswith("module.final"):
                        del state_dict[key]
            state_dict = {k[7:]: w for k, w in state_dict.items()}
            model.load_state_dict(state_dict, strict=False)
            if not args.from_zero:
                start_epoch = checkpoint['epoch']
                if not args.zero_score:
                    dice_best = checkpoint.get('dice_best', 0)
                    f_score_best = checkpoint.get('f_score_best', 0)
            print("=> loaded checkpoint '{}' (epoch {})"
                  .format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))
    if args.from_zero:
        start_epoch = 0
    current_epoch = start_epoch

    if conf['fp16']:
        model, optimizer = amp.initialize(model, optimizer,
                                          opt_level=args.opt_level,
                                          loss_scale='dynamic')

    if args.distributed:
        model = DistributedDataParallel(model, delay_allreduce=True)
    else:
        model = DataParallel(model).cuda()
    for epoch in range(start_epoch, conf['optimizer']['schedule']['epochs']):
        if train_sampler:
            train_sampler.set_epoch(epoch)

        model.module.train()
        train_epoch(current_epoch, loss_functions, model, optimizer, scheduler, train_data_loader, summary_writer, conf,
                    args.local_rank)

        model = model.eval()
        if args.local_rank == 0:
            torch.save({
                'epoch': current_epoch + 1,
                'state_dict': model.state_dict(),
                'dice_best': dice_best,
                'f_score_best': f_score_best,
            }, args.output_dir + '/' + snapshot_name + "_last")
            if epoch % args.test_every == 0:
                preds_dir = os.path.join(args.predictions, snapshot_name)
                dice_best, f_score_best = evaluate_val(args, val_data_loader, dice_best, f_score_best,
                                         model,
                                         snapshot_name=snapshot_name,
                                         current_epoch=current_epoch,
                                         optimizer=optimizer, summary_writer=summary_writer,
                                         predictions_dir=preds_dir, data_dir=args.data_dir, visualizer_path=args.visualizer_path)
        current_epoch += 1

def evaluate_val(args, data_val, dice_best, f_score_best, model, snapshot_name, current_epoch, optimizer, summary_writer,
                 predictions_dir, data_dir, visualizer_path):
    print("Test phase")
    model = model.eval()
    dice, f_score = validate(model, data_loader=data_val, predictions_dir=predictions_dir, data_dir=data_dir,
                             visualizer_path=visualizer_path)
    if args.local_rank == 0:
        summary_writer.add_scalar('val/dice', float(dice), global_step=current_epoch)
        summary_writer.add_scalar('val/F1', f_score, global_step=current_epoch)
        if dice > dice_best:
            print("Dice improved from {} to {}".format(dice_best, dice))
            if args.output_dir is not None:
                torch.save({
                    'epoch': current_epoch + 1,
                    'state_dict': model.state_dict(),
                    'dice_best': dice,
                    'dice': dice,
                    'f_score': f_score,
                    'f_score_best': f_score_best,
                }, args.output_dir + snapshot_name + "_best_dice")
            dice_best = dice
        if f_score > f_score_best:
            print("F1 improved from {} to {}".format(f_score_best, f_score))
            if args.output_dir is not None:
                torch.save({
                    'epoch': current_epoch + 1,
                    'state_dict': model.state_dict(),
                    'dice_best': dice_best,
                    'dice': dice,
                    'f_score': f_score,
                    'f_score_best': f_score,
                }, args.output_dir + snapshot_name + "_best_f1")
            f_score_best = f_score
        torch.save({
            'epoch': current_epoch + 1,
            'state_dict': model.state_dict(),
            'dice_best': dice_best,
            'f_score_best': f_score_best,
            'dice': dice,
            'f_score': f_score,
        }, args.output_dir + snapshot_name + "_last")
        print("dice: {}, dice_best: {}".format(dice, dice_best))
    return dice_best, f_score_best

def validate(net, data_loader, predictions_dir, data_dir, visualizer_path):
    os.makedirs(predictions_dir, exist_ok=True)
    preds_dir = predictions_dir + "/predictions"
    os.makedirs(preds_dir, exist_ok=True)
    dices = []
    with torch.no_grad():
        for sample in tqdm(data_loader):
            imgs = sample["image"].cuda().float()
            mask = sample["mask"].cuda().float()

            output = net(imgs)
            binary_pred = torch.sigmoid(output)

            for i in range(output.shape[0]):
                d = dice_round(binary_pred[:, 0:1, :], mask[:, 0:1, ...], t=0.5).item()
                dices.append(d)
                cv2.imwrite(os.path.join(preds_dir, sample["img_name"][i] + ".png"),
                            (np.moveaxis(binary_pred[i].cpu().numpy(), 0, -1)[..., :3] * 255))
    f_score = calculate_metrics(fold_dir=predictions_dir,
                                visualizer_path=visualizer_path,
                                truth_csv=os.path.join(data_dir,
                                                       "SummaryData/SN6_Train_AOI_11_Rotterdam_Buildings.csv"),
                                img_dir=os.path.join(data_dir, "SAR-Intensity/"),
                                sar_orientations_csv=os.path.join(data_dir, "SummaryData/SAR_orientations.txt"))
    return np.mean(dices), f_score

def train_epoch(current_epoch, loss_functions, model, optimizer, scheduler, train_data_loader, summary_writer, conf,
                local_rank):
    losses = AverageMeter()
    dices = AverageMeter()
    iterator = tqdm(train_data_loader)
    model.train()
    if conf["optimizer"]["schedule"]["mode"] == "epoch":
        scheduler.step(current_epoch)
    for i, sample in enumerate(iterator):
        imgs = sample["image"].cuda()
        masks = sample["mask"].cuda().float()
        out_mask = model(imgs)
        with torch.no_grad():
            pred = torch.sigmoid(out_mask)
            d = dice_round(pred[:, 0:1, ...], masks[:, 0:1, ...], t=0.5).item()
        dices.update(d, imgs.size(0))

        mask_loss = loss_functions["mask_loss"](out_mask, masks)
        loss = mask_loss
        losses.update(loss.item(), imgs.size(0))
        iterator.set_description(
            "epoch: {}; lr {:.7f}; Loss ({loss.avg:.4f}); dice ({dice.avg:.4f}); ".format(
                current_epoch, scheduler.get_lr()[-1], loss=losses, dice=dices))
        optimizer.zero_grad()
        if conf['fp16']:
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            loss.backward()
        torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), 1)
        optimizer.step()
        torch.cuda.synchronize()

        if conf["optimizer"]["schedule"]["mode"] in ("step", "poly"):
            scheduler.step(i + current_epoch * len(train_data_loader))

    if local_rank == 0:
        for idx, param_group in enumerate(optimizer.param_groups):
            lr = param_group['lr']
            summary_writer.add_scalar('group{}/lr'.format(idx), float(lr), global_step=current_epoch)
        summary_writer.add_scalar('train/loss', float(losses.avg), global_step=current_epoch)

if __name__ == '__main__':
    main()


from collections import OrderedDict
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
import math

import gym
from gym import spaces
import numpy as np

import torch
from torch import nn
import torch.nn.functional as F
from torch import distributions

import pytorch_lightning as pl

from lightning_baselines3.off_policy_models import OffPolicyModel
from lightning_baselines3.common.type_aliases import GymEnv

class SAC(OffPolicyModel):
    """
    Soft Actor-Critic (SAC)
    Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor,
    This implementation borrows code from original implementation (https://github.com/haarnoja/sac)
    from OpenAI Spinning Up (https://github.com/openai/spinningup), from the softlearning repo
    (https://github.com/rail-berkeley/softlearning/)
    and from Stable Baselines (https://github.com/hill-a/stable-baselines)
    Paper: https://arxiv.org/abs/1801.01290
    Introduction to SAC: https://spinningup.openai.com/en/latest/algorithms/sac.html

    Note: we use double q target and not value target as discussed
    in https://github.com/hill-a/stable-baselines/issues/270

    :param env: The environment to learn from
        (if registered in Gym, can be str. Can be None for loading trained models)
    :param eval_env: The environment to evaluate on, must not be vectorised/parallelrised
        (if registered in Gym, can be str. Can be None for loading trained models)
    :param batch_size: Minibatch size for each gradient update
    :param buffer_length: length of the replay buffer
    :param warmup_length: how many steps of the model to collect transitions for before learning starts
    :param train_freq: Update the model every ``train_freq`` steps. Set to `-1` to disable.
    :param episodes_per_rollout: Update the model every ``episodes_per_rollout`` episodes.
        Note that this cannot be used at the same time as ``train_freq``. Set to `-1` to disable.
    :param num_rollouts: Number of rollouts to do per PyTorch Lightning epoch. This does not affect any training dynamic,
        just how often we evaluate the model since evaluation happens at the end of each Lightning epoch
    :param gradient_steps: How many gradient steps to do after each rollout
    :param target_update_interval: How many environment steps to wait between updating the target Q network
    :param num_eval_episodes: The number of episodes to evaluate for at the end of a PyTorch Lightning epoch
    :param gamma: the discount factor
    :param entropy_coef: Entropy regularization coefficient. (Equivalent to
        inverse of reward scale in the original SAC paper.)  Controlling exploration/exploitation trade-off.
        Set it to 'auto' to learn it automatically (and 'auto_0.1' for using 0.1 as initial value)
    :param target_entropy: target entropy when learning ``ent_coef`` (``ent_coef = 'auto'``)
    :param use_sde: Whether to use generalized State Dependent Exploration (gSDE)
        instead of action noise exploration (default: False)
    :param sde_sample_freq: Sample a new noise matrix every n steps when using gSDE
        Default: -1 (only sample at the beginning of the rollout)
    :param use_sde_at_warmup: Whether to use gSDE instead of uniform sampling
        during the warm up phase (before learning starts)
    :param squashed_actions: Whether the actions are squashed between [-1, 1] and need to be unsquashed
    :param verbose: The verbosity level: 0 none, 1 training information, 2 debug (default: 0)
    :param seed: Seed for the pseudo random generators
    """

    def __init__(
        self,
        env: Union[GymEnv, str],
        eval_env: Union[GymEnv, str],
        batch_size: int = 256,
        buffer_length: int = int(1e6),
        warmup_length: int = 100,
        train_freq: int = 1,
        episodes_per_rollout: int = -1,
        num_rollouts: int = 1000,
        gradient_steps: int = 1,
        target_update_interval: int = 1,
        num_eval_episodes: int = 10,
        gamma: float = 0.99,
        entropy_coef: Union[str, float] = "auto",
        target_entropy: Union[str, float] = "auto",
        use_sde: bool = False,
        sde_sample_freq: int = -1,
        use_sde_at_warmup: bool = False,
        squashed_actions: bool = True,
        verbose: int = 0,
        seed: Optional[int] = None,
    ):
        self.entropy_coef = entropy_coef
        self.target_entropy = target_entropy
        self.target_update_interval = target_update_interval

        super(SAC, self).__init__(
            env=env,
            eval_env=eval_env,
            batch_size=batch_size,
            buffer_length=buffer_length,
            warmup_length=warmup_length,
            train_freq=train_freq,
            episodes_per_rollout=episodes_per_rollout,
            num_rollouts=num_rollouts,
            gradient_steps=gradient_steps,
            num_eval_episodes=num_eval_episodes,
            gamma=gamma,
            verbose=verbose,
            squashed_actions=squashed_actions,
            seed=seed,
            use_sde=use_sde,
            use_sde_at_warmup=use_sde_at_warmup,
            sde_sample_freq=sde_sample_freq
        )

        assert isinstance(self.action_space, spaces.Box), "SAC only supports environments with Box action spaces"

        # We need manual optimization for this
        self.automatic_optimization = False

    def reset(self):
        """
        Resets the environment and automatic entropy
        """
        super(SAC, self).reset()
        # Target entropy is used when learning the entropy coefficient
        if self.target_entropy == "auto":
            # automatically set target entropy if needed
            self.target_entropy = -np.prod(self.env.action_space.shape).astype(np.float32)
        else:
            # Force conversion
            # this will also throw an error for unexpected string
            self.target_entropy = float(self.target_entropy)

        if isinstance(self.entropy_coef, str):
            if not hasattr(self, 'log_entropy_coef'):
                assert self.entropy_coef.startswith("auto")
                # Default initial value of entropy_coef when learned
                init_value = 1.0
                if "_" in self.entropy_coef:
                    init_value = float(self.entropy_coef.split("_")[1])
                    assert init_value > 0.0, "The initial value of entropy_coef must be greater than 0"

                # Note: we optimize the log of the entropy coeff which is slightly different from the paper
                # as discussed in https://github.com/rail-berkeley/softlearning/issues/37
                self.log_entropy_coef = torch.log(torch.ones(1, device=self.device) * init_value)
                self.log_entropy_coef = nn.Parameter(self.log_entropy_coef.requires_grad_(True))
                self.entropy_coef_optimizer = torch.optim.Adam([self.log_entropy_coef], lr=3e-4)
        else:
            # I know this isn't very efficient but it makes the code cleaner
            # and it's only one extra operation
            self.log_entropy_coef = torch.log(float(self.entropy_coef))

    def forward_actor(self, obs: torch.Tensor) -> torch.Tensor:
        """
        Runs the actor network.
        Override this function with your own.

        :param obs: The input observations
        :return: The deterministic action of the actor
        """
        raise NotImplementedError

    def forward_critics(
        self, obs: torch.Tensor, action: torch.Tensor
        )-> Tuple[torch.Tensor, ...]:
        """
        Runs the all critic networks.
        Override this function with your own.

        :param obs: The input observations
        :param action: The input actions
        :return: The output Q values of the critic networks in the form of a list
        """
        raise NotImplementedError

    def forward_critic_targets(
        self, obs: torch.Tensor, action: torch.Tensor
        )-> Tuple[torch.Tensor, ...]:
        """
        Runs the all target critic networks.
        Override this function with your own.

        :param obs: The input observations
        :param action: The input actions
        :return: The output Q values of the critic networks in the form of a list
        """
        raise NotImplementedError

    def update_targets(self) -> None:
        """
        Function to update the target critic networks periodically.
        Override this function with your own.
        """
        raise NotImplementedError

    def training_step(self, batch, batch_idx, optimizer_idx):
        """
        Specifies the update step for SAC. Override this if you wish to modify the SAC algorithm
        """
        # We need to sample because `log_std` may have changed between two gradient steps
        if self.num_timesteps < self.warmup_length:
            return

        opt_critic, opt_actor = self.optimizers(use_pl_optimizer=True)

        # Action by the current actor for the sampled state
        dist = self.forward_actor(batch.observations)
        actions = dist.rsample()
        log_probs = dist.log_prob(actions)

        log_entropy_coef = torch.clamp(self.log_entropy_coef, -10, 5)
        entropy_coef = torch.exp(log_entropy_coef)
        if hasattr(self, 'entropy_coef_optimizer'):
            # Important: detach the variable from the graph
            # so we don't change it with other losses
            # see https://github.com/rail-berkeley/softlearning/issues/60
            entropy_coef = entropy_coef.detach()
            entropy_coef_loss = -(log_entropy_coef * (log_probs + self.target_entropy).detach()).mean()
            self.log('entropy_coef_loss', entropy_coef_loss, on_step=True, prog_bar=True, logger=True)

        self.log('entropy_coef', entropy_coef, on_step=True, prog_bar=False, logger=True)

        # Optimize entropy coefficient, also called
        # entropy temperature or alpha in the paper
        if hasattr(self, 'entropy_coef_optimizer'):
            self.entropy_coef_optimizer.zero_grad()
            self.manual_backward(entropy_coef_loss, self.entropy_coef_optimizer)
            self.entropy_coef_optimizer.step()

        with torch.no_grad():
            # Select action according to policy
            next_dist = self.forward_actor(batch.next_observations)
            next_actions = next_dist.rsample()
            next_log_probs = next_dist.log_prob(next_actions)

            # Compute the target Q value: min over all critics targets
            targets = self.forward_critic_targets(batch.next_observations, next_actions)
            target_q = torch.minimum(*targets)
            # add entropy term
            target_q = target_q - entropy_coef * next_log_probs[..., None]
            # td error + entropy term
            target_q = batch.rewards + (1 - batch.dones) * self.gamma * target_q

        # Get current Q estimates for each critic network
        # using action from the replay buffer
        current_q_estimates = self.forward_critics(batch.observations, batch.actions)

        # Compute critic loss
        critic_loss = torch.stack([F.mse_loss(current_q, target_q) for current_q in current_q_estimates])
        critic_loss = torch.mean(critic_loss)
        self.log('critic_loss', critic_loss, on_step=True, prog_bar=True, logger=True)

        # Optimize the critic
        opt_critic.zero_grad()
        self.manual_backward(critic_loss, opt_critic)
        opt_critic.step()

        # Compute actor loss
        # Alternative: actor_loss = torch.mean(log_prob - qf1_pi)
        # Mean over all critic networks
        q_values_pi = self.forward_critics(batch.observations, actions)
        min_qf_pi = torch.minimum(*q_values_pi)
        actor_loss = (entropy_coef * log_probs[..., None] - min_qf_pi).mean()
        self.log('actor_loss', actor_loss, on_step=True, prog_bar=True, logger=True)

        # Optimize the actor
        opt_actor.zero_grad()
        self.manual_backward(actor_loss, opt_actor)
        opt_actor.step()

        # Update target networks
        if batch_idx % self.target_update_interval == 0:
            self.update_targets()


from __future__ import print_function
import datetime
import os
import time
import sys
import numpy as np
import torch
import torch.utils.data
from torch.utils.data.dataloader import default_collate
from torch import nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms

import utils

from datasets.msr import MSRAction3D
import models.sequence_classification as Models

def train_one_epoch(model, criterion, optimizer, lr_scheduler, data_loader, device, epoch, print_freq):
    model.train()
    metric_logger = utils.MetricLogger(delimiter="  ")
    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value}'))
    metric_logger.add_meter('clips/s', utils.SmoothedValue(window_size=10, fmt='{value:.3f}'))

    header = 'Epoch: [{}]'.format(epoch)
    for clip, target, _ in metric_logger.log_every(data_loader, print_freq, header):
        start_time = time.time()
        clip, target = clip.to(device), target.to(device)
        output = model(clip)
        loss = criterion(output, target)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
        batch_size = clip.shape[0]
        metric_logger.update(loss=loss.item(), lr=optimizer.param_groups[0]["lr"])
        metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
        metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
        metric_logger.meters['clips/s'].update(batch_size / (time.time() - start_time))
        lr_scheduler.step()
        sys.stdout.flush()

def evaluate(model, criterion, data_loader, device):
    model.eval()
    metric_logger = utils.MetricLogger(delimiter="  ")
    header = 'Test:'
    video_prob = {}
    video_label = {}
    with torch.no_grad():
        for clip, target, video_idx in metric_logger.log_every(data_loader, 100, header):
            clip = clip.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True)
            output = model(clip)
            loss = criterion(output, target)

            acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
            prob = F.softmax(input=output, dim=1)

            # FIXME need to take into account that the datasets
            # could have been padded in distributed setup
            batch_size = clip.shape[0]
            target = target.cpu().numpy()
            video_idx = video_idx.cpu().numpy()
            prob = prob.cpu().numpy()
            for i in range(0, batch_size):
                idx = video_idx[i]
                if idx in video_prob:
                    video_prob[idx] += prob[i]
                else:
                    video_prob[idx] = prob[i]
                    video_label[idx] = target[i]
            metric_logger.update(loss=loss.item())
            metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
            metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
    # gather the stats from all processes
    metric_logger.synchronize_between_processes()

    print(' * Clip Acc@1 {top1.global_avg:.3f} Clip Acc@5 {top5.global_avg:.3f}'.format(top1=metric_logger.acc1, top5=metric_logger.acc5))

    # video level prediction
    video_pred = {k: np.argmax(v) for k, v in video_prob.items()}
    pred_correct = [video_pred[k]==video_label[k] for k in video_pred]
    total_acc = np.mean(pred_correct)

    class_count = [0] * data_loader.dataset.num_classes
    class_correct = [0] * data_loader.dataset.num_classes

    for k, v in video_pred.items():
        label = video_label[k]
        class_count[label] += 1
        class_correct[label] += (v==label)
    class_acc = [c/float(s) for c, s in zip(class_correct, class_count)]

    print(' * Video Acc@1 %f'%total_acc)
    print(' * Class Acc@1 %s'%str(class_acc))

    return total_acc

def main(args):

    if args.output_dir:
        utils.mkdir(args.output_dir)

    print(args)
    print("torch version: ", torch.__version__)
    print("torchvision version: ", torchvision.__version__)

    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    device = torch.device('cuda')

    # Data loading code
    print("Loading data")

    st = time.time()

    dataset = MSRAction3D(
            root=args.data_path,
            frames_per_clip=args.clip_len,
            step_between_clips=1,
            num_points=args.num_points,
            train=True
    )

    dataset_test = MSRAction3D(
            root=args.data_path,
            frames_per_clip=args.clip_len,
            step_between_clips=1,
            num_points=args.num_points,
            train=False
    )

    print("Creating data loaders")

    data_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True)

    data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=args.batch_size, num_workers=args.workers, pin_memory=True)

    print("Creating model")
    Model = getattr(Models, args.model)
    model = Model(radius=args.radius, nsamples=args.nsamples, num_classes=dataset.num_classes)
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)
    model.to(device)

    criterion = nn.CrossEntropyLoss()

    lr = args.lr
    optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=args.momentum, weight_decay=args.weight_decay)

    # convert scheduler to be per iteration, not per epoch, for warmup that lasts
    # between different epochs
    warmup_iters = args.lr_warmup_epochs * len(data_loader)
    lr_milestones = [len(data_loader) * m for m in args.lr_milestones]
    lr_scheduler = utils.WarmupMultiStepLR(optimizer, milestones=lr_milestones, gamma=args.lr_gamma, warmup_iters=warmup_iters, warmup_factor=1e-5)

    model_without_ddp = model

    if args.resume:
        checkpoint = torch.load(args.resume, map_location='cpu')
        model_without_ddp.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
        args.start_epoch = checkpoint['epoch'] + 1

    print("Start training")
    start_time = time.time()
    acc = 0
    for epoch in range(args.start_epoch, args.epochs):
        train_one_epoch(model, criterion, optimizer, lr_scheduler, data_loader, device, epoch, args.print_freq)

        acc = max(acc, evaluate(model, criterion, data_loader_test, device=device))

        if args.output_dir:
            checkpoint = {
                'model': model_without_ddp.state_dict(),
                'optimizer': optimizer.state_dict(),
                'lr_scheduler': lr_scheduler.state_dict(),
                'epoch': epoch,
                'args': args}
            utils.save_on_master(
                checkpoint,
                os.path.join(args.output_dir, 'model_{}.pth'.format(epoch)))
            utils.save_on_master(
                checkpoint,
                os.path.join(args.output_dir, 'checkpoint.pth'))

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))
    print('Accuracy {}'.format(acc))

def parse_args():
    import argparse
    parser = argparse.ArgumentParser(description='PSTNet Training')

    parser.add_argument('--data-path', default='data/MSR-Action3D', type=str, help='dataset')
    parser.add_argument('--seed', default=0, type=int, help='random seed')
    parser.add_argument('--model', default='MSRAction', type=str, help='model')
    parser.add_argument('--radius', default=0.5, type=float, help='radius for the ball query')
    parser.add_argument('--nsamples', default=9, type=int, help='number of neighbors for the ball query')
    parser.add_argument('--clip-len', default=16, type=int, metavar='N', help='number of frames per clip')
    parser.add_argument('--num-points', default=2048, type=int, metavar='N', help='number of points per frame')
    parser.add_argument('-b', '--batch-size', default=16, type=int)
    parser.add_argument('--epochs', default=35, type=int, metavar='N', help='number of total epochs to run')
    parser.add_argument('-j', '--workers', default=10, type=int, metavar='N', help='number of data loading workers (default: 16)')
    parser.add_argument('--lr', default=0.01, type=float, help='initial learning rate')
    parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum')
    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)', dest='weight_decay')
    parser.add_argument('--lr-milestones', nargs='+', default=[20, 30], type=int, help='decrease lr on milestones')
    parser.add_argument('--lr-gamma', default=0.1, type=float, help='decrease lr by a factor of lr-gamma')
    parser.add_argument('--lr-warmup-epochs', default=10, type=int, help='number of warmup epochs')
    parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
    parser.add_argument('--output-dir', default='output/MAR-Action3D', type=str, help='path where to save')
    parser.add_argument('--resume', default='', help='resume from checkpoint')
    parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='start epoch')

    args = parser.parse_args()

    return args

if __name__ == "__main__":
    args = parse_args()
    main(args)


'''
Module contains all of the classes that represent Machine Learning models
that are within `Tang et al. 2016 paper \
<https://aclanthology.info/papers/C16-1311/c16-1311>`_:

1. :py:class:`bella.models.tdlstm.LSTM` -- LSTM model.
2. :py:class:`bella.models.tdlstm.TDLSTM` -- TDLSTM model.
3. :py:class:`bella.models.tdlstm.TCLSTM` -- TCLSTM model.
'''

from typing import Dict, Callable, Any, List, Union, Tuple

import numpy as np
import keras
from keras import models, optimizers, layers
from keras.utils import to_categorical

import bella
from bella.models.base import KerasModel
from bella.contexts import context
from bella.neural_pooling import matrix_median

class LSTM(KerasModel):
    '''
    Attributes:

    1. pad_size -- The max number of tokens to use per sequence. If -1
       use the text sequence in the training data that has the most tokens as
       the pad size.
    2. embedding_layer_kwargs -- Keyword arguments to pass to the embedding
       layer which is a :py:class:`keras.layers.Embedding` object. Can be
       None if no parameters to pass.
    3. lstm_layer_kwargs -- Keyword arguments to pass to the lstm layer(s)
       which is a :py:class:`keras.layers.LSTM` object. Can be
       None if no parameters to pass.
    4. dense_layer_kwargs -- Keyword arguments to pass to the dense (final
       layer) which is a :py:class:`keras.layers.Dense` object. Can be
       None if no parameters to pass.

    Methods:

    1. model_parameters -- Returns a dictionary containing the attributes of
       the class instance, the parameters to give to the class constructior to
       re-create this instance, and the class itself.
    2. create_training_text -- Converts the training and validation data into a
       format that the keras model can take as input.
    3. create_training_y -- Converts the training and validation target values
       from a vector of class lables into a matrix of binary values. of shape
       [n_samples, n_classes].
    4. keras_model -- The model that represents this class. This is a single
       forward LSTM.
    '''

    @classmethod
    def name(cls) -> str:
        return 'LSTM'

    def __repr__(self) -> str:
        '''
        Name of the machine learning model.
        '''
        return self.name()

    def __init__(self, tokeniser: Callable[[str], List[str]],
                 embeddings: 'bella.word_vectors.WordVectors',
                 reproducible: Union[int, None] = None, pad_size: int = -1,
                 lower: bool = True, patience: int = 10,
                 batch_size: int = 32, epochs: int = 300,
                 embedding_layer_kwargs: Dict[str, Any] = None,
                 lstm_layer_kwargs: Dict[str, Any] = None,
                 dense_layer_kwargs: Dict[str, Any] = None,
                 optimiser: 'keras.optimizers.Optimizer' = optimizers.SGD,
                 optimiser_params: Union[Dict[str, Any], None] = None
                 ) -> None:
        '''
        :param tokeniser: Tokeniser to be used e.g. :py:meth:`str.split`.
        :param embeddings: Embedding (Word vectors) to be used e.g.
                           :py:class:`bella.word_vectors.SSWE`
        :param reproducible: Whether to be reproducible. If None then it is
                             quicker to run. Else provide a `int` that
                             will represent the random seed value.
        :param pad_size: The max number of tokens to use per sequence. If -1
                         use the text sequence in the training data that has
                         the most tokens as the pad size.
        :param lower: Whether to lower case the words being processed.
        :param patience: Number of epochs with no improvement before training
                         is stopped.
        :param batch_size: Number of samples per gradient update.
        :param epochs: Number of times to train over the entire training set
                       before stopping. If patience is set, then it may
                       stop before reaching the number of epochs specified
                       here.
        :param embedding_layer_kwargs: Keyword arguments to pass to the
                                       embedding layer which is a
                                       :py:class:`keras.layers.Embedding`
                                       object. If no parameters to pass leave
                                       as None.
        :param lstm_layer_kwargs: Keyword arguments to pass to the lstm
                                  layer(s) which is a
                                  :py:class:`keras.layers.LSTM` object. If no
                                  parameters to pass leave as None.
        :param dense_layer_kwargs: Keyword arguments to pass to the dense
                                   (final layer) which is a
                                   :py:class:`keras.layers.Dense` object. If no
                                   parameters to pass leave as None.
        :param optimiser: Optimiser to be used accepts any
                          `keras optimiser <https://keras.io/optimizers/>`_.
                          Default is :py:class:`keras.optimizers.SGD`
        :param optimiser_params: Parameters for the optimiser. If None uses
                                 default optimiser parameters.
        '''

        self.tokeniser = tokeniser
        self.embeddings = embeddings
        self.reproducible = reproducible
        self.pad_size = pad_size
        self.test_pad_size = 0
        self.lower = lower
        self.patience = patience
        self.batch_size = batch_size
        self.epochs = epochs

        self.optimiser_params = optimiser_params
        if optimiser_params is None:
            self.optimiser_params = {}

        self.embedding_layer_kwargs = embedding_layer_kwargs
        if embedding_layer_kwargs is None:
            self.embedding_layer_kwargs = {}
        self.lstm_layer_kwargs = lstm_layer_kwargs
        if lstm_layer_kwargs is None:
            self.lstm_layer_kwargs = {}
        self.dense_layer_kwargs = dense_layer_kwargs
        if dense_layer_kwargs is None:
            self.dense_layer_kwargs = {}
        self.optimiser = optimiser
        self.model = None
        self.fitted = False

    def model_parameters(self) -> Dict[str, Any]:
        '''
        Returns a dictionary containing the attributes of the class instance,
        the parameters to give to the class constructior to re-create this
        instance, and the class itself.

        This is used by the :py:meth:`save` method so that the instance can
        be re-created when loaded by the :py:meth:`load` method.
        '''

        class_params = {'tokeniser': self.tokeniser,
                        'embeddings': self.embeddings,
                        'reproducible': self.reproducible,
                        'pad_size': self.pad_size,
                        'lower': self.lower,
                        'patience': self.patience,
                        'batch_size': self.batch_size, 'epochs': self.epochs,
                        'embedding_layer_kwargs': self.embedding_layer_kwargs,
                        'lstm_layer_kwargs': self.lstm_layer_kwargs,
                        'dense_layer_kwargs': self.dense_layer_kwargs,
                        'optimiser': self.optimiser,
                        'optimiser_params': self.optimiser_params}
        class_attrs = {'test_pad_size': self.test_pad_size}
        return {'class': self.__class__, 'class_attrs': class_attrs,
                'class_params': class_params}

    def _pre_process(self, data_dicts: List[Dict[str, str]],
                     training: bool = False) -> np.ndarray:
        '''
        Converts the text in the data_dicts into a matrix of shape
        [n_samples, pad_size] where each integer in the matrix represents
        the word embedding lookup. This is then used as input into the
        keras model.

        The text from the data_dicts are converted by the
        :py:meth:`process_text` method.

        :param data_dicts: A list of dictonaries that contains a `text` field.
        :param training: Whether the text should be processed for training or
                         for prediction. prediction = False, training = True
        :return: The output of :py:meth:`process_text` method.
        '''
        text_data = [data['text'] for data in data_dicts]
        if training:
            pad_data = self.process_text(text_data, self.pad_size)
            self.test_pad_size, sequence_data = pad_data
            return sequence_data
        _, sequence_data = self.process_text(text_data, self.test_pad_size)
        return sequence_data

    def create_training_y(self, train_y: np.ndarray, validation_y: np.ndarray,
                          ) -> Tuple[np.ndarray, np.ndarray]:
        '''
        Converts the training and validation target values from a vector of
        class lables into a matrix of binary values of shape [n_samples,
        n_classes].

        To convert the vector of classes to a matrix we the
        :py:func:`keras.utils.to_categorical` function.

        :param train_y: Vector of class labels, shape = [n_samples]
        :param validation_y: Vector of class labels, shape = [n_samples]
        :return: A tuple of length two containing the train and validation
                 matrices respectively. The shape of each matrix is:
                 [n_samples, n_classes]
        '''
        train_y = to_categorical(train_y).astype(np.float32)
        validation_y = to_categorical(validation_y).astype(np.float32)
        return train_y, validation_y

    def create_training_text(self, train_data: List[Dict[str, str]],
                             validation_data: List[Dict[str, str]]
                             ) -> Tuple[np.ndarray, np.ndarray]:
        '''
        Converts the training and validation data into a format that the keras
        model can take as input.

        :param train_data: Data to be trained on. Which is a list of
                           dictionaries where each dictionary has a `text`
                           field containing text.
        :param validation_data: Data to evaluate the model at training time.
                                Which is a list of dictionaries where each
                                dictionary has a `text` field containing text.
        :return: A tuple of length two containing the train and validation
                 input that are both the output of :py:meth:`_pre_process`
        '''
        train_sequence = self._pre_process(train_data, training=True)
        val_sequence = self._pre_process(validation_data, training=False)
        return train_sequence, val_sequence

    def keras_model(self, num_classes: int) -> 'keras.models.Model':
        '''
        The model that represents this class. This is a single forward LSTM.

        :param num_classes: Number of classes to predict.
        :return: Forward LSTM keras model.
        '''
        # Embeddings
        embedding_matrix = self.embeddings.embedding_matrix
        vocab_size, vector_size = embedding_matrix.shape

        embedding_layer_kwargs = self.embedding_layer_kwargs
        embedding_layer_trainable = True
        if 'trainable' in embedding_layer_kwargs:
            embedding_layer_trainable = embedding_layer_kwargs.pop('trainable')

        lstm_layer_kwargs = self.lstm_layer_kwargs
        lstm_dimension = vector_size
        if 'cell' in self.lstm_layer_kwargs:
            lstm_dimension = lstm_layer_kwargs.pop('cell')

        dense_layer_kwargs = self.dense_layer_kwargs
        # output_activation = 'softmax' if num_classes > 2 else ''
        # Model layers
        input_layer = layers.Input(shape=(self.test_pad_size,),
                                   name='text_input')
        embedding_layer = layers\
                          .Embedding(input_dim=vocab_size,
                                     output_dim=vector_size,
                                     input_length=self.test_pad_size,
                                     trainable=embedding_layer_trainable,
                                     weights=[embedding_matrix],
                                     name='embedding_layer',
                                     **embedding_layer_kwargs
                                     )(input_layer)
        lstm_layer = layers.LSTM(lstm_dimension,
                                 name='lstm_layer',
                                 **lstm_layer_kwargs)(embedding_layer)
        prediction_layer = layers.Dense(num_classes, activation='softmax',
                                        name='output',
                                        **dense_layer_kwargs)(lstm_layer)

        return models.Model(inputs=input_layer, outputs=prediction_layer)

    @property
    def pad_size(self) -> int:
        '''
        pad_size attribute

        :return: The pad_size used in the model
        '''

        return self._pad_size

    @pad_size.setter
    def pad_size(self, value: int) -> None:
        '''
        Sets the pad_size attribute

        :param value: The value to assign to the pad_size attribute
        '''

        self.fitted = False
        self._pad_size = value

    @property
    def embedding_layer_kwargs(self) -> Dict[str, Any]:
        '''
        embedding_layer_kwargs attribute

        :return: The embedding_layer_kwargs used in the model
        '''

        return self._embedding_layer_kwargs

    @embedding_layer_kwargs.setter
    def embedding_layer_kwargs(self, value: Dict[str, Any]) -> None:
        '''
        Sets the embedding_layer_kwargs attribute

        :param value: The value to assign to the embedding_layer_kwargs
                      attribute
        '''

        self.fitted = False
        self._embedding_layer_kwargs = value

    @property
    def lstm_layer_kwargs(self) -> Dict[str, Any]:
        '''
        lstm_layer_kwargs attribute

        :return: The lstm_layer_kwargs used in the model
        '''

        return self._lstm_layer_kwargs

    @lstm_layer_kwargs.setter
    def lstm_layer_kwargs(self, value: Dict[str, Any]) -> None:
        '''
        Sets the lstm_layer_kwargs attribute

        :param value: The value to assign to the lstm_layer_kwargs
                      attribute
        '''

        self.fitted = False
        self._lstm_layer_kwargs = value

    @property
    def dense_layer_kwargs(self) -> Dict[str, Any]:
        '''
        dense_layer_kwargs attribute

        :return: The dense_layer_kwargs used in the model
        '''

        return self._dense_layer_kwargs

    @dense_layer_kwargs.setter
    def dense_layer_kwargs(self, value: Dict[str, Any]) -> None:
        '''
        Sets the dense_layer_kwargs attribute

        :param value: The value to assign to the dense_layer_kwargs
                      attribute
        '''

        self.fitted = False
        self._dense_layer_kwargs = value

class TDLSTM(LSTM):
    '''
    Attributes:

    1. include_target -- Wheather to include the target in the LSTM
       representations.
    '''

    @classmethod
    def name(cls) -> str:
        return 'TDLSTM'

    def __repr__(self) -> str:
        '''
        Name of the machine learning model.
        '''
        return self.name()

    def __init__(self, tokeniser: Callable[[str], List[str]],
                 embeddings: 'bella.word_vectors.WordVectors',
                 reproducible: Union[int, None] = None, pad_size: int = -1,
                 lower: bool = True, patience: int = 10,
                 batch_size: int = 32, epochs: int = 300,
                 embedding_layer_kwargs: Dict[str, Any] = None,
                 lstm_layer_kwargs: Dict[str, Any] = None,
                 dense_layer_kwargs: Dict[str, Any] = None,
                 optimiser: 'keras.optimizers.Optimizer' = optimizers.SGD,
                 optimiser_params: Union[Dict[str, Any], None] = None,
                 include_target: bool = True) -> None:
        '''
        :param tokeniser: Tokeniser to be used e.g. :py:meth:`str.split`.
        :param embeddings: Embedding (Word vectors) to be used e.g.
                           :py:class:`bella.word_vectors.SSWE`
        :param reproducible: Whether to be reproducible. If None then it is
                             but quicker to run. Else provide a `int` that
                             will represent the random seed value.
        :param pad_size: The max number of tokens to use per sequence. If -1
                         use the text sequence in the training data that has
                         the most tokens as the pad size.
        :param lower: Whether to lower case the words being processed.
        :param patience: Number of epochs with no improvement before training
                         is stopped.
        :param batch_size: Number of samples per gradient update.
        :param epochs: Number of times to train over the entire training set
                       before stopping. If patience is set, then it may
                       stop before reaching the number of epochs specified
                       here.
        :param embedding_layer_kwargs: Keyword arguments to pass to the
                                       embedding layer which is a
                                       :py:class:`keras.layers.Embedding`
                                       object. If no parameters to pass leave
                                       as None.
        :param lstm_layer_kwargs: Keyword arguments to pass to the lstm
                                  layer(s) which is a
                                  :py:class:`keras.layers.LSTM` object. If no
                                  parameters to pass leave as None.
        :param dense_layer_kwargs: Keyword arguments to pass to the dense
                                   (final layer) which is a
                                   :py:class:`keras.layers.Dense` object. If no
                                   parameters to pass leave as None.
        :param optimiser: Optimiser to be used accepts any
                          `keras optimiser <https://keras.io/optimizers/>`_.
                          Default is :py:class:`keras.optimizers.SGD`
        :param optimiser_params: Parameters for the optimiser. If None uses
                                 default optimiser parameters.
        :param include_target: Wheather to include the target in the LSTM
                               representations.
        '''

        super().__init__(tokeniser, embeddings, reproducible, pad_size, lower,
                         patience, batch_size, epochs, embedding_layer_kwargs,
                         lstm_layer_kwargs, dense_layer_kwargs, optimiser,
                         optimiser_params)

        self.left_pad_size = pad_size
        self.left_test_pad_size = 0
        self.right_pad_size = pad_size
        self.right_test_pad_size = 0
        self.include_target = include_target

    def model_parameters(self) -> Dict[str, Any]:
        '''
        Returns a dictionary containing the attributes of the class instance,
        the parameters to give to the class constructior to re-create this
        instance, and the class itself.

        This is used by the :py:meth:`save` method so that the instance can
        be re-created when loaded by the :py:meth:`load` method.
        '''

        attributes = super().model_parameters()
        class_attrs = {'left_test_pad_size': self.left_test_pad_size,
                       'right_test_pad_size': self.right_test_pad_size}
        attributes['class_attrs'] = class_attrs

        class_params = attributes['class_params']
        class_params['include_target'] = self.include_target
        attributes['class_params'] = class_params
        return attributes

    def _pre_process(self, data_dicts: List[Dict[str, Any]],
                     training: bool = False) -> List[np.ndarray]:
        '''
        Converts the text in the data_dicts into a List of size two
        representing the left and right context of the target word
        respectively. Each List is made up of a matrix of of integers
        representing the text as their embedding lookups. These two Lists
        are the inputs into the keras model.

        Two find the left and right contexts it uses the `spans` field of
        the dictionaries in the `data_dicts`. The `spans` field is a list of
        Tuples where each Tuple represents a occurence of the Target, each
        Tuple contains the index of the starting and ending character index
        (Expects the List to be of size 1 as there should be only one target
        per target sample. This case is not True for the
        `Dong et al. <https://aclanthology.info/papers/P14-2009/p14-2009>`_
        dataset therefore it only takes the first target instance in the
        sentence as the target).

        The texts are converted into integers using the
        :py:meth:`process_text` method.

        :param data_dicts: A list of dictonaries that contains a `text` and
                           `spans` field.
        :param training: Whether the text should be processed for training or
                         for prediction. prediction = False, training = True
        :return: A list of two contaning the left and right context of
                 the target both represented by the output of
                 :py:meth:`process_text` method.
        '''

        def context_texts(context_data_dicts: List[Dict[str, Any]]
                          ) -> Tuple[List[str], List[str]]:
            '''
            :param context_data_dicts: A list of dictonaries that contains a
                                       `text` and `spans` field.
            :return: A list of the left and right text contexts for all the
                     dictionaries.
            '''
            # Context returns all of the left and right context occurrences
            # therefore if a target is mentioned Twice and are associated then
            # for a single text two left and right occurrences are returned.
            # Thus these are a list of lists we therefore chose only the
            # first mentioned target as the paper linked to this method does
            # not specify which they used.
            left_texts = [context(data, 'left', inc_target=self.include_target)
                          for data in context_data_dicts]
            right_texts = [context(data, 'right',
                                   inc_target=self.include_target)
                           for data in context_data_dicts]
            left_texts = [texts[0] for texts in left_texts]
            right_texts = [texts[0] for texts in right_texts]
            return left_texts, right_texts

        # Convert from a sequence of dictionaries into texts and then integers
        # that represent the tokens in the text within the embedding space.

        # Get left and right contexts
        left_text, right_text = context_texts(data_dicts)
        if training:
            left_pad_sequence = self.process_text(left_text,
                                                  self.left_pad_size)
            self.left_test_pad_size, left_sequence = left_pad_sequence

            right_pad_sequence = self.process_text(right_text,
                                                   self.right_pad_size,
                                                   padding='post',
                                                   truncate='post')
            self.right_test_pad_size, right_sequence = right_pad_sequence
            return left_sequence, right_sequence

        left_pad_sequence = self.process_text(left_text,
                                              self.left_test_pad_size)
        _, left_sequence = left_pad_sequence

        right_pad_sequence = self.process_text(right_text,
                                               self.right_test_pad_size,
                                               padding='post',
                                               truncate='post')
        _, right_sequence = right_pad_sequence
        return [left_sequence, right_sequence]

    def create_training_text(self, train_data: List[Dict[str, Any]],
                             validation_data: List[Dict[str, Any]]
                             ) -> Tuple[List[np.ndarray],
                                        List[np.ndarray]]:
        '''
        Converts the training and validation data into a format that the keras
        model can take as input.

        :param train_data: Data to be trained on. Which is a list of
                           dictionaries where each dictionary has a `text`
                           field containing text and a field `spans` containing
                           a list of Tuples where each Tuple represents a
                           occurence of the Target, each Tuple contains the
                           index of the starting and ending character index
                           (Expects the List to be of size 1 as there should
                           be only one target per target sample. This case is
                           not True for the
                           `Dong et al. <https://aclanthology.info/papers/P14-\
                           2009/p14-2009>`_ dataset therefore it only takes
                           the first target instance in the sentence as the
                           target).
        :param validation_data: Data to evaluate the model at training time.
                                Expects the same data as the `train_data`
                                parameter.
        :return: A tuple of length two containing the train and validation
                 input that are both the output of :py:meth:`_pre_process`
        '''

        return super().create_training_text(train_data, validation_data)

    def keras_model(self, num_classes: int) -> 'keras.models.Model':
        '''
        The model that represents this class. This is a custom combination
        of two LSTMs.

        :param num_classes: Number of classes to predict.
        :return: Two LSTMs, one forward from the left context and the other
                 backward from the right context. The output of the two are
                 concatenated and are input to the output layer.
        '''
        # Embeddings
        embedding_matrix = self.embeddings.embedding_matrix
        vocab_size, vector_size = embedding_matrix.shape

        embedding_layer_kwargs = self.embedding_layer_kwargs
        embedding_layer_trainable = True
        if 'trainable' in embedding_layer_kwargs:
            embedding_layer_trainable = embedding_layer_kwargs.pop('trainable')

        lstm_layer_kwargs = self.lstm_layer_kwargs
        lstm_dimension = vector_size
        if 'cell' in self.lstm_layer_kwargs:
            lstm_dimension = lstm_layer_kwargs.pop('cell')

        dense_layer_kwargs = self.dense_layer_kwargs
        # Model layers
        # Left LSTM
        left_input = layers.Input(shape=(self.left_test_pad_size,),
                                  name='left_text_input')
        left_embedding_layer = layers\
                               .Embedding(input_dim=vocab_size,
                                          output_dim=vector_size,
                                          input_length=self.left_test_pad_size,
                                          trainable=embedding_layer_trainable,
                                          weights=[embedding_matrix],
                                          name='left_embedding_layer',
                                          **embedding_layer_kwargs
                                          )(left_input)
        left_lstm_layer = layers.LSTM(lstm_dimension,
                                      name='left_lstm_layer',
                                      **lstm_layer_kwargs
                                      )(left_embedding_layer)
        # Right LSTM
        right_input = layers.Input(shape=(self.right_test_pad_size,),
                                   name='right_text_input')
        right_embedding_layer = layers\
                                .Embedding(input_dim=vocab_size,
                                           output_dim=vector_size,
                                           input_length=self.right_test_pad_size,
                                           trainable=embedding_layer_trainable,
                                           weights=[embedding_matrix],
                                           name='right_embedding_layer',
                                           **embedding_layer_kwargs
                                           )(right_input)
        right_lstm_layer = layers.LSTM(lstm_dimension,
                                       name='right_lstm_layer',
                                       go_backwards=True,
                                       **lstm_layer_kwargs
                                       )(right_embedding_layer)
        # Merge the outputs of the left and right LSTMs
        merge_layer = layers.concatenate([left_lstm_layer, right_lstm_layer],
                                         name='left_right_lstm_merge')
        predictions = layers.Dense(num_classes, activation='softmax',
                                   name='output',
                                   **dense_layer_kwargs)(merge_layer)

        return models.Model(inputs=[left_input, right_input],
                            outputs=predictions)

    @property
    def include_target(self) -> bool:
        '''
        include_target attribute

        :return: The include_target used in the model
        '''

        return self._include_target

    @include_target.setter
    def include_target(self, value: bool) -> None:
        '''
        Sets the include_target attribute

        :param value: The value to assign to the include_target attribute
        '''

        self.fitted = False
        self._include_target = value

class TCLSTM(TDLSTM):

    @classmethod
    def name(cls) -> str:
        return 'TCLSTM'

    def __repr__(self) -> str:
        '''
        Name of the machine learning model.
        '''
        return self.name()

    def _pre_process(self, data_dicts: List[Dict[str, Any]],
                     training: bool = False) -> List[np.ndarray]:
        '''
        Converts the text in the data_dicts into a list of size four
        representing the left context, left targets, right context and
        right targets. Where the contexts come are the same as those from
        TDLSTM :py:meth:`bella.models.tdlstm.TDLSTM._pre_process` method.

        The targets are a matrix of size [word_embedding_dimension, pad_size]
        and each vector in the matrix is the word embedding representation
        of the target word. If the target word is made up of multiple words
        it is then the average of the words vector representation (we use the
        median as the average). Both the contexts and the target matrix are
        used as input into the keras model.

        The texts are converted into integers using the
        :py:meth:`process_text` method.

        :param data_dicts: A list of dictonaries that contains a `text` and
                           `spans` field.
        :param training: Whether the text should be processed for training or
                         for prediction. prediction = False, training = True
        :return: A list of four contaning the left context, left vectors,
                 right context, and right vectors.
        '''
        def context_median_targets(pad_size: int):
            '''
            :param pad_size: The number of timesteps within the LSTM
            :return: Matrix of size [word_embedding_dimension, pad_size] where
                     each word embedding represents the target word or if
                     multiple words make up the target the word embedding is
                     the median of the words embeddings.
            '''
            vector_size = self.embeddings.vector_size
            target_matrix = np.zeros((len(data_dicts),
                                      pad_size, vector_size))
            for index, data in enumerate(data_dicts):
                target_vectors = []
                target_words = data['target'].split()
                for target_word in target_words:
                    if self.lower:
                        target_word = target_word.lower()
                    target_embedding = self.embeddings\
                                           .lookup_vector(target_word)
                    target_vectors.append(target_embedding)
                target_vectors = np.vstack(target_vectors)
                median_target_vector = matrix_median(target_vectors)
                median_vectors = np.repeat(median_target_vector, pad_size,
                                           axis=0)
                target_matrix[index] = median_vectors
            return target_matrix

        sequences = super()._pre_process(data_dicts, training=training)
        left_sequence, right_sequence = sequences
        left_target_vectors = context_median_targets(self.left_test_pad_size)
        right_target_vectors = context_median_targets(self.right_test_pad_size)
        return [left_sequence, left_target_vectors,
                right_sequence, right_target_vectors]

    def create_training_text(self, train_data: List[Dict[str, Any]],
                             validation_data: List[Dict[str, Any]]
                             ) -> Tuple[List[np.ndarray],
                                        List[np.ndarray]]:
        '''
        Converts the training and validation data into a format that the keras
        model can take as input.

        :param train_data: See :py:meth:`bella.models.tdlstm.\
                           TDLSTM.create_training_text` `train_data`
                           parameter.
        :param validation_data: See :py:meth:`bella.models.tdlstm.\
                                TDLSTM.create_training_text` `validation_data`
                                parameter.
        :return: A tuple of length two containing the train and validation
                 input that are both the output of :py:meth:`_pre_process`
        '''

        return super().create_training_text(train_data, validation_data)

    def keras_model(self, num_classes: int) -> 'keras.models.Model':
        '''
        The model that represents this class. This is the same as the
        :py:meth:`bella.models.tdlstm.TDLSTM.keras_model` model, however
        the words in before inputting into the LSTM are concatenated with
        the word embedding of the target. If the target is more than one word
        then the word embedding of the target is the average (median in our
        case) embeddings of the target words.

        :param num_classes: Number of classes to predict.
        :return: Two LSTMs one forward from the left context and the other
                 backward from the right context taking into account the
                 target vector embedding.
        '''
        # Embeddings
        embedding_matrix = self.embeddings.embedding_matrix
        vocab_size, vector_size = embedding_matrix.shape

        embedding_layer_kwargs = self.embedding_layer_kwargs
        embedding_layer_trainable = True
        if 'trainable' in embedding_layer_kwargs:
            embedding_layer_trainable = embedding_layer_kwargs.pop('trainable')

        lstm_layer_kwargs = self.lstm_layer_kwargs
        # Double the vector size as we have to take into consideration the
        # concatenated target vector
        lstm_dimension = vector_size * 2
        if 'cell' in self.lstm_layer_kwargs:
            lstm_dimension = lstm_layer_kwargs.pop('cell')

        dense_layer_kwargs = self.dense_layer_kwargs
        # Model layers
        # Left LSTM
        left_input = layers.Input(shape=(self.left_test_pad_size,),
                                  name='left_text_input')
        left_embedding_layer = layers\
                               .Embedding(input_dim=vocab_size,
                                          output_dim=vector_size,
                                          input_length=self.left_test_pad_size,
                                          trainable=embedding_layer_trainable,
                                          weights=[embedding_matrix],
                                          name='left_embedding_layer',
                                          **embedding_layer_kwargs
                                          )(left_input)
        left_target_input = layers.Input(shape=(self.left_test_pad_size,
                                                vector_size),
                                         name='left_target')
        left_text_target = layers.concatenate([left_embedding_layer,
                                               left_target_input],
                                              name='left_text_target')
        left_lstm_layer = layers.LSTM(lstm_dimension,
                                      name='left_lstm_layer',
                                      **lstm_layer_kwargs
                                      )(left_text_target)
        # Right LSTM
        right_input = layers.Input(shape=(self.right_test_pad_size,),
                                   name='right_text_input')
        right_embedding_layer = layers\
                                .Embedding(input_dim=vocab_size,
                                           output_dim=vector_size,
                                           input_length=self.right_test_pad_size,
                                           trainable=embedding_layer_trainable,
                                           weights=[embedding_matrix],
                                           name='right_embedding_layer',
                                           **embedding_layer_kwargs
                                           )(right_input)
        right_target_input = layers.Input(shape=(self.right_test_pad_size,
                                                 vector_size),
                                          name='right_target')
        right_text_target = layers.concatenate([right_embedding_layer,
                                                right_target_input],
                                               name='right_text_target')
        right_lstm_layer = layers.LSTM(lstm_dimension,
                                       name='right_lstm_layer',
                                       go_backwards=True,
                                       **lstm_layer_kwargs
                                       )(right_text_target)
        # Merge the outputs of the left and right LSTMs
        merge_layer = layers.concatenate([left_lstm_layer, right_lstm_layer],
                                         name='left_right_lstm_merge')
        predictions = layers.Dense(num_classes, activation='softmax',
                                   name='output',
                                   **dense_layer_kwargs)(merge_layer)

        input_layers = [left_input, left_target_input,
                        right_input, right_target_input]
        return models.Model(inputs=input_layers, outputs=predictions)


from screenshot import screenshot
from config import green_bar, green, check_button, raise_button, bet_button, call_button, fold_button, all_in_button
from config import raise_button_sum, fold_button_sum, check_button_sum, call_button_sum, bet_button_sum, all_in_button_sum
import cv2
import numpy as np

def player_turn_check():
    green_bar_np = np.array(screenshot(green_bar))
    green_bar_gray = cv2.cvtColor(green_bar_np, cv2.COLOR_RGB2BGR)
    if green in green_bar_gray:
        return True
    else:
        return False

def available_moves_check():
    moves = []
    check_button_np = np.array(screenshot(check_button))
    check_button_gray = cv2.cvtColor(check_button_np, cv2.COLOR_RGB2GRAY)
    # cv2.imwrite('check %i.png' %check_button_gray.sum(), check_button_gray)
    if check_button_gray.sum() == check_button_sum:
        moves.append("check")

    raise_button_np = np.array(screenshot(raise_button))
    raise_button_gray = cv2.cvtColor(raise_button_np, cv2.COLOR_RGB2GRAY)
    # cv2.imwrite('raise %i.png' %raise_button_gray.sum(), raise_button_gray)
    if raise_button_gray.sum() == raise_button_sum:
        moves.append("raise")

    fold_button_np = np.array(screenshot(fold_button))
    fold_button_gray = cv2.cvtColor(fold_button_np, cv2.COLOR_RGB2GRAY)
    # cv2.imwrite('fold %i.png' %fold_button_gray.sum(), fold_button_gray)
    if fold_button_gray.sum() == fold_button_sum:
        moves.append("fold")

    call_button_np = np.array(screenshot(call_button))
    call_button_gray = cv2.cvtColor(call_button_np, cv2.COLOR_RGB2GRAY)
    # cv2.imwrite('call %i.png' %call_button_gray.sum(), call_button_gray)
    if call_button_gray.sum() == call_button_sum:
        moves.append("call")

    bet_button_np = np.array(screenshot(bet_button))
    bet_button_gray = cv2.cvtColor(bet_button_np, cv2.COLOR_RGB2GRAY)
    # cv2.imwrite('bet %i.png' %bet_button_gray.sum(), bet_button_gray)
    if bet_button_gray.sum() == bet_button_sum:
        moves.append("bet")

    all_in_button_np = np.array(screenshot(all_in_button))
    all_in_button_gray = cv2.cvtColor(all_in_button_np, cv2.COLOR_RGB2GRAY)
    # cv2.imwrite('all_in %i.png' %all_in_button_gray.sum(), all_in_button_gray)
    if all_in_button_gray.sum() == all_in_button_sum:
        moves.append("all_in")

    return moves