#
# General Electricity sector Decarbonization Model (GEDM)
# Copyright (C) 2020 Cheng-Ta Chu.
# Licensed under the MIT License (see LICENSE file).
#
# Module note:
# Functions for import settings at country level
#

from copy import deepcopy
from numpy import genfromtxt, vstack

import cls_misc

def get_CountryTechCostAssump(instance):
    ''' import technical and cost assumptions at country level '''
    # technology assumptions at country level
    # object inherit from base processes

    # copy the default instance setting
    for country in instance.lsCountry:
        country.lsProcessAssump = deepcopy(instance.lsProcessDefObjs)

    ### import technical assumptions

    techno_economic = []
    file_process = "../Input/3_process/02_Tech_coal.csv"
    techno_economic = genfromtxt(file_process, dtype = str, skip_header=0, delimiter=',')

    file_process = "../Input/3_process/02_Tech_gas.csv"
    techno_economic = vstack( (techno_economic, genfromtxt(file_process, \
                                    dtype = str, skip_header=1, delimiter=',')) )

    file_process = "../Input/3_process/02_Tech_other.csv"
    techno_economic = vstack( (techno_economic, genfromtxt(file_process, \
                                    dtype = str, skip_header=1, delimiter=',')) )

    file_process = "../Input/3_process/02_Tech_renewable.csv"
    techno_economic = vstack( (techno_economic, genfromtxt(file_process, \
                                    dtype = str, skip_header=1, delimiter=',')) )

    # year step period in the original setting
    iOriginYear_OYS = techno_economic[0, 4:]

    for country in instance.lsCountry:
        for process in country.lsProcessAssump:
            sProcessName = process.sProcessName
            sRegionGroup = country.sTechCostGroup

            bProcessFlag = False
            for assumption in techno_economic:
                if sProcessName == assumption[0]:
                    bProcessFlag = True

                    if sRegionGroup == assumption[1]:

                        if assumption[2] == "UnitCapacity":
                            process.iUnitCapacity = assumption[4]

                        elif assumption[2] == "Gross_Eff":
                            process.fGrossEff_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                        elif assumption[2] == "MinLoad":
                            process.fMinLoad_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                        elif assumption[2] == "RampRate":
                            process.fRampRate_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                        elif assumption[2] == "EquAvailFactor":
                            process.fEquAvailFactor_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                        elif assumption[2] == "AuxiliaryCon":
                            process.fAuxiliaryCon_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                        elif assumption[2] == "CaptureRate":
                            process.fCaptureRate_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                        elif assumption[2] == "Duration":
                            process.fDuration_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                else:
                    if bProcessFlag == True:
                        break

    ### import cost assumptions

    techno_economic = []
    file_process = "../Input/3_process/03_Cost_coal.csv"
    techno_economic = genfromtxt(file_process, dtype = str, skip_header=1, delimiter=',')

    file_process = "../Input/3_process/03_Cost_gas.csv"
    techno_economic = vstack( (techno_economic, genfromtxt(file_process, \
                                dtype = str, skip_header=1, delimiter=',')) )

    file_process = "../Input/3_process/03_Cost_other.csv"
    techno_economic = vstack( (techno_economic, genfromtxt(file_process, \
                                dtype = str, skip_header=1, delimiter=',')) )

    file_process = "../Input/3_process/03_Cost_renewable.csv"
    techno_economic = vstack( (techno_economic, genfromtxt(file_process, \
                                dtype = str, skip_header=1, delimiter=',')) )

    for country in instance.lsCountry:
        for process in country.lsProcessAssump:
            sProcessName = process.sProcessName
            sRegionGroup = country.sTechCostGroup

            bProcessFlag = False
            for assumption in techno_economic:
                if sProcessName == assumption[0]:
                    bProcessFlag = True

                    if sRegionGroup == assumption[1]:

                        if assumption[2] == "CAPEX":
                            process.fCAPEX_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                        elif assumption[2] == "OPEX":
                            process.fOPEX_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                        elif assumption[2] == "varOPEX":
                            process.fVarOPEX_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                            assumption[4:], instance.iAllYearSteps_YS)

                        elif assumption[2] == "Lifetime":
                            process.fLifetime = float(assumption[4])

                else:
                    if bProcessFlag == True:
                        break

        # annualized fix cost per MW (# M.USD / yr.MW)
        for process in country.lsProcessAssump:
            for iYS in range( 0, len(instance.iAllYearSteps_YS)):

                ### cost assumptions
                fCapacity = 1                                 # MW
                fCapitalCost = process.fCAPEX_YS[iYS] * fCapacity * 1000 # USD/KW * MW * 1000 = USD
                fYearOMCost = process.fOPEX_YS[iYS] * fCapacity * 1000  # USD/KW * MW * 1000 = USD
                fDiscountRate = process.fDiscount
                iPlantLife = process.fLifetime
                # fCapitalRecoveyFactor = (D*(1+D)^L) / ( ((1+D)^L)-1 )
                fCapitalRecoveyFactor =  ( fDiscountRate * ((1+fDiscountRate)**iPlantLife)) \
                    / ( ((1+fDiscountRate)**iPlantLife) - 1 )
                process.fAnnualCapex[iYS] = fCapitalCost * fCapitalRecoveyFactor / 1000000  # M.USD/year
                process.fAnnualFixedCost[iYS] = process.fAnnualCapex[iYS] \
                    + (fYearOMCost / 1000000)  # M.USD/year

    return

def get_CountryCommodityAssump(instance):
    ''' import commodity assumptions '''
    # import commodity assumption at country level
    # commodity assumptions are all applied at country level

    file_commodity = "../Input/2_commodity/03_CommodityCost.csv"
    comm_assump = genfromtxt(file_commodity, dtype = str, skip_header=0, delimiter=',')

    file_commodity = "../Input/2_commodity/02_CostGroup.csv"
    comm_group = genfromtxt(file_commodity, dtype = str, skip_header=1, delimiter=',')

    # year step period in the original setting
    iOriginYear_OYS = comm_assump[0, 3:]

    for country in instance.lsCountry:

        # look up cost group
        CG_oil = "NPS"
        CG_gas = "NPS_EU"
        CG_coal = "NPS_EU"
        CG_uranium = "NPS"
        CG_biomass = "NPS_EU"
        for row in comm_group:
            if row[1] == country.sCountry:
                CG_oil = row[2]
                CG_gas = row[3]
                CG_coal = row[4]
                CG_uranium = row[5]
                CG_biomass = row[6]
                break

        country.lsCommodity = deepcopy(instance.lsCommodity)

        for commodity in country.lsCommodity:

            Category = commodity.sCategory
            CostGroup = ""
            if Category == "oil":
                CostGroup = CG_oil
            elif Category == "gas":
                CostGroup = CG_gas
            elif Category == "coal":
                CostGroup = CG_coal
            elif Category == "uranium":
                CostGroup = CG_uranium
            elif Category == "biomass":
                CostGroup = CG_biomass

            for assumption in comm_assump:
                if Category == assumption[0] and CostGroup == assumption[1]:
                    commodity.fFuelPrice_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                        assumption[3:], instance.iAllYearSteps_YS)

    return

def get_CountryPolicyAndTechAssump(instance):
    ''' import country policy and target assumptions '''

    ### carbon cost
    file_data = "../Input/4_policy/02_CarbonCost.csv"
    carbon_cost = genfromtxt(file_data, dtype = str, skip_header=0, delimiter=',')

    # year step period in the original setting
    iOriginYear_OYS = carbon_cost[0, 2:]

    for country in instance.lsCountry:
        for row_cost in carbon_cost:
            if row_cost[0] == country.sCarbonCostGroup:
                country.fCarbonCost_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                        row_cost[2:], instance.iAllYearSteps_YS)
                break

    ### ancillary service requirement
    file_data = "../Input/4_policy/03_AncillaryService.csv"
    AS_requirement = genfromtxt(file_data, dtype = str, skip_header=0, delimiter=',')

    # year period steps in the original setting
    iOriginYear_OYS = AS_requirement[0, 4:]

    for country in instance.lsCountry:
        for row_AS in AS_requirement:
            if row_AS[0] == country.sAncillaryGroup and row_AS[1] == "Primary":
                country.fASFirst_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                        row_AS[4:], instance.iAllYearSteps_YS)
                country.fASFirst_YS = [ cell/100 for cell in country.fASFirst_YS ]

            elif row_AS[0] == country.sAncillaryGroup and row_AS[1] == "Secondary":
                country.fASSecond_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                        row_AS[4:], instance.iAllYearSteps_YS)
                country.fASSecond_YS = [ cell/100 for cell in country.fASSecond_YS ]

            elif row_AS[0] == country.sAncillaryGroup and row_AS[1] == "Tertiary":
                country.fASThird_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                        row_AS[4:], instance.iAllYearSteps_YS)
                country.fASThird_YS = [ cell/100 for cell in country.fASThird_YS ]

    ### dispatchable process new build limit at specific period (CCS, Nuke)
    file_data = "../Input/4_policy/04_FixedNewBuild_dispatch.csv"
    PolicyDevLimit = genfromtxt(file_data, dtype = str, skip_header=0, delimiter=',')

    # year step period in the original setting
    iOriginYear_OYS = PolicyDevLimit[0, 3:]

    for country in instance.lsCountry:
        for process in country.lsProcessAssump:
            for row in PolicyDevLimit:
                if row[1] == country.sCountry and row[2] == process.sProcessName:

                    for iYear in instance.iAllYearSteps_YS:
                        for indCol, sValue in enumerate(PolicyDevLimit[0]):
                            if indCol >= 3 and row[indCol] != "":
                                if iYear == int(sValue):
                                    # restriction only applied when it matches the year step period
                                    process.dicProcDispFixedNewBuild[str(iYear)] = float(row[indCol])
                                    break
                                elif iYear < int(sValue):
                                    # if the year step does not match, set 0
                                    process.dicProcDispFixedNewBuild[str(iYear)] = 0
                                    break
                    break

    '''
    ### renewable target
    file_data = "../Input/4_policy/01_RE_Target.csv"
    RETarget = genfromtxt(file_data, dtype = str, skip_header=1, delimiter=',')
    for country in instance.lsCountry:
        for row in RETarget:
            if row[0] == country.sCountry:
                country.lsRenewableTarget = row[1:]
    '''

    return

def get_CountryBaseYearData(instance):
    ''' import country base year data assumptions '''
    # only "own use in industry" and "system loss" are used now

    file_data = "../Input/5_calibration/01_BaseYear_Generation.csv"
    Base_Data = genfromtxt(file_data, dtype = str, skip_header=1, delimiter=',')

    for country in instance.lsCountry:

        # generation by type
        gen_coal = -1
        gen_gas = -1
        gen_oil = -1
        gen_nuke = -1
        gen_bio = -1

        # conversion eff by type
        eff_coal = -1
        eff_gas = -1
        eff_oil = -1
        eff_bio = -1

        # flow
        flow_total_gen = -1
        flow_import = -1
        flow_export = -1
        flow_ownuse = -1
        flow_ownindustry = -1
        flow_loss = -1

        for row in Base_Data:
            if row[0] == country.sCountry:
                gen_coal = float(row[2])
                gen_gas = float(row[3])
                gen_oil = float(row[4])
                gen_nuke = float(row[5])
                gen_bio = float(row[6])
                # conversion eff by type
                eff_coal = float(row[7])
                eff_gas = float(row[8])
                eff_oil = float(row[9])
                eff_bio = float(row[11])
                # flow
                flow_total_gen = float(row[12])
                flow_import = float(row[17])
                flow_export = - float(row[18])  # turn into positive value
                flow_ownuse = float(row[20]) / flow_total_gen
                flow_ownindustry = float(row[22]) / flow_total_gen
                flow_loss = float(row[23]) / flow_total_gen
                break

        country.dicBaseYearData["gen_coal"] = gen_coal
        country.dicBaseYearData["gen_gas"] = gen_gas
        country.dicBaseYearData["gen_oil"] = gen_oil
        country.dicBaseYearData["gen_nuke"] = gen_nuke
        country.dicBaseYearData["gen_bio"] = gen_bio
        country.dicBaseYearData["eff_coal"] = eff_coal
        country.dicBaseYearData["eff_gas"] = eff_gas
        country.dicBaseYearData["eff_oil"] = eff_oil
        country.dicBaseYearData["eff_bio"] = eff_bio
        country.dicBaseYearData["flow_total_gen"] = flow_total_gen
        country.dicBaseYearData["flow_import"] = flow_import
        country.dicBaseYearData["flow_export"] = flow_export
        country.dicBaseYearData["flow_ownuse"] = flow_ownuse
        country.dicBaseYearData["flow_ownindustry"] = flow_ownindustry
        country.dicBaseYearData["flow_loss"] = flow_loss

    return

def get_CountryRenInstallLimit(instance):
    ''' import settings on renewable installation limits '''

    ### limits applied on if the setting year matches the time period
    # minimum installation capacity
    file_cdata = "../Input/8_VRE/01_RE_Min_Add.csv"
    zone_ReMin = genfromtxt(file_cdata, dtype = str, skip_header=0, delimiter=',')
    header = zone_ReMin[0]
    for objCountry in instance.lsCountry:
        for row_Data in zone_ReMin:
            if row_Data[0] == objCountry.sCountry:
                tech_type = row_Data[1]
                for ind_col, tech_cap in enumerate(row_Data):
                    if ind_col > 1:
                        if (tech_cap != "") and ( int(header[ind_col]) in instance.iAllYearSteps_YS):
                            objCountry.dicRenMinInstall[tech_type,header[ind_col]] = float(tech_cap)

    # maximum installation capacity
    file_cdata = "../Input/8_VRE/02_RE_Max_Add.csv"
    zone_ReMax = genfromtxt(file_cdata, dtype = str, skip_header=0, delimiter=',')
    header = zone_ReMax[0]
    for objCountry in instance.lsCountry:
        for row_Data in zone_ReMax:
            if row_Data[0] == objCountry.sCountry:
                tech_type = row_Data[1]
                for ind_col, tech_cap in enumerate(row_Data):
                    if ind_col > 1:
                        if (tech_cap != "") and ( int(header[ind_col]) in instance.iAllYearSteps_YS):
                            objCountry.dicRenMaxInstall[tech_type,header[ind_col]] = float(tech_cap)

    return

def get_CountryHydroCapLimit(instance):
    ''' import hydropower maximum capacity limits '''

    file_cdata = "../Input/8_VRE/03_RE_Hydro_Max_Cap.csv"
    country_HydroCap = genfromtxt(file_cdata, dtype = str, skip_header=0, delimiter=',')

    for objCountry in instance.lsCountry:

        for row_country in country_HydroCap:
            if row_country[0] == objCountry.sCountry:
                if row_country[2] != "":
                    objCountry.fTotalHydroCapLimit = float(row_country[2])
                else:
                    objCountry.fTotalHydroCapLimit = 0
                break
    return

def get_CountryBiomassLimit(instance):
    ''' import the projection of available biomass for power production '''

    file_cdata = "../Input/8_VRE/04_RE_Biomass_Max_Pot.csv"
    country_Data = genfromtxt(file_cdata, dtype = str, skip_header=0, delimiter=',')

    # year step period in the original setting
    iOriginYear_OYS = country_Data[0, 1:]

    for objCountry in instance.lsCountry:
        for row_country in country_Data:
            if row_country[0] == objCountry.sCountry:
                objCountry.fBiomassLimit_YS = ConvertAssumpYS(iOriginYear_OYS, \
                                        row_country[1:], instance.iAllYearSteps_YS)
                break
    return

def get_MarketTransmission(instance, objMarket):
    ''' import transmission link configurations in a market'''
    # create link objects

    ### terrestrial zones
    file_conn = "../Input/6_transmission/01_Zone_conn_base.csv"
    zone_conn = genfromtxt(file_conn, dtype = str, skip_header=1, delimiter=',')

    for ind_row, zone_ter in enumerate(zone_conn):

        for zone in objMarket.lsZone:

            if zone.sZone == zone_ter[2] or zone.sZone == zone_ter[3]:
                FromID = ""
                ToID = ""
                # get zoneID
                for objZone in objMarket.lsZone:
                    if zone_ter[2] == objZone.sZone:
                        FromID = objZone.sZoneID
                        break
                for objZone in objMarket.lsZone:
                    if zone_ter[3] == objZone.sZone:
                        ToID = objZone.sZoneID
                        break

                # could connect to a zone which is not in the market
                if ToID == "":
                    ToID = "ET" + str(ind_row)
                if FromID == "":
                    FromID = "ET" + str(ind_row)

                # check existing object
                bDupicate = False
                for objTrans in objMarket.lsTrans:
                    if objTrans.sFrom == FromID and objTrans.sTo == ToID:
                        bDupicate == True
                        break

                if bDupicate == False:
                    fDist = float(zone_ter[6])
                    fBaseCap = float(zone_ter[7])
                    # if the base year connection exist, set a minimal 100MW NTC
                    # the value will be updated
                    if zone_ter[4] == "1":
                        fBaseCap = max( fBaseCap, 100 )
                    objTrans = cls_misc.Transmission( From=FromID, To=ToID, \
                        Dist=fDist, Conn2015=zone_ter[4], BaseCap=fBaseCap )
                    objMarket.lsTrans.append(objTrans)

                break

    ### offshore zones
    file_conn = "../Input/6_transmission/02_Zone_conn_base_offshore.csv"
    zone_conn = genfromtxt(file_conn, dtype = str, skip_header=1, delimiter=',')

    for zone in objMarket.lsZoneOffs:

        bProcessFlag = False
        for zone_off in zone_conn:

            if zone_off[2] == zone.sZone:
                bProcessFlag = True
                FromID = ""
                ToID = ""
                # get zoneID
                for objZone in objMarket.lsZoneOffs:
                    if zone_off[2] == objZone.sZone:
                        FromID = objZone.sZoneID
                        break
                for objZone in objMarket.lsZone:
                    if zone_off[3] == objZone.sZone:
                        ToID = objZone.sZoneID
                        break

                fDist = float(zone_off[5])
                fBaseCap = float(zone_off[6])
                objTrans = cls_misc.Transmission( From=FromID, To=ToID, \
                    Dist=fDist, Conn2015=zone_off[4], BaseCap=fBaseCap )
                objMarket.lsTrans_off.append(objTrans)
            else:
                if bProcessFlag == True:
                    break

    return

def get_TransTechEcoAssump(instance, objMarket):
    ''' import general techno-economic assumptions of transmission '''

    file_commodity = "../Input/3_process/04_Transmission.csv"
    comm_assump = genfromtxt(file_commodity, dtype = str, skip_header=0, delimiter=',')

    # year step period in the original setting
    iOriginYear_OYS = comm_assump[0, 3:]

    for ind_row, data_row in enumerate(comm_assump):

        if data_row[0] == "DC_Line_Loss":
            objMarket.lsDCLineLoss = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

        elif data_row[0] == "DC_Conv_Loss":
            objMarket.lsDCConvLoss = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

        elif data_row[0] == "AC_Line_Loss":
            objMarket.lsACLineLoss = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

        elif data_row[0] == "DC_CAPEX":
            objMarket.lsDCCapex = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

        elif data_row[0] == "DC_OPEX":
            objMarket.lsDCOpex = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

        elif data_row[0] == "DC_CAPEX_conv":
            objMarket.lsDCCapexConv = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

        elif data_row[0] == "DC_OPEX_conv":
            objMarket.lsDCOpexConv = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

        elif data_row[0] == "AC_CAPEX":
            objMarket.lsACCapex = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

        elif data_row[0] == "AC_OPEX":
            objMarket.lsACOpex = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

        elif data_row[0] == "CRF":
            objMarket.lsCRF = ConvertAssumpYS(iOriginYear_OYS, \
                                    data_row[3:], instance.iAllYearSteps_YS)

    return

def ConvertAssumpYS(sOriginYear_OYS, sAssumption_OYS, iAllYearSteps_YS):
    ''' convert the data in setting files into defined periods steps '''
    # the year step setting in the CSV files may not match the time period setting
    # using intepolation

    iAssumptions_YS = []

    for iYear in iAllYearSteps_YS:
        for ind_OriYear, sOriYear in enumerate(sOriginYear_OYS):
            if iYear <= int(sOriYear):
                iAssumptions_YS.append( float(sAssumption_OYS[ind_OriYear]) )
                break
            elif ind_OriYear == (len(sOriginYear_OYS) - 1):
                # last element in the original assumption list
                iAssumptions_YS.append( float(sAssumption_OYS[ind_OriYear]) )
                break
            elif iYear < int(sOriginYear_OYS[ind_OriYear+1]):
                # intepolate for the value
                iInterval = abs(int(sOriginYear_OYS[ind_OriYear + 1]) \
                                - int(sOriginYear_OYS[ind_OriYear]))
                iInterval_1 = int(sOriginYear_OYS[ind_OriYear+1]) - iYear
                iInterval_2 = iYear - int(sOriginYear_OYS[ind_OriYear])
                fIntpValue = ( float(sAssumption_OYS[ind_OriYear])*iInterval_1 \
                              + float(sAssumption_OYS[ind_OriYear+1])*iInterval_2 ) / iInterval
                iAssumptions_YS.append(fIntpValue)
                break

    return iAssumptions_YS


# Load libraries
import argparse
import Lib_ as cf
import numpy as np
import os
import scanpy as sc
import warnings
from numba.core.errors import NumbaPerformanceWarning
warnings.simplefilter('ignore', category=NumbaPerformanceWarning) # Fix this with update if possible

def main(AnnData, output_folder, **params):
	# Set warning settings
	warnings.simplefilter(action = 'ignore', category = FutureWarning)
	sc.settings.verbosity = 0 # Only print errors

	# Check if output directory already exist else make a new directory.
	if os.path.isdir('{}/{}'.format(path, output_folder)) == True:
		print('Output folder is already present.')
		print('Continuing may overright files in folder...')
	else:
		os.mkdir('{}/{}'.format(path, output_folder))

	# Load datafile.
	print('Loading {}.'.format(AnnData))
	adata = sc.read(AnnData)
	print('File details: {}'.format(adata))
	# Check whether raw attribute exists.
	try:
		adata.raw.X
	except AttributeError:
		print('adata.raw.X does not exist.')
		print('Please set attribute: https://anndata.readthedocs.io/en/latest/anndata.AnnData.raw.html')
		print('adata.raw.X should contain raw count matrix (not logarithmized).')
		exit()

	print('Adjusting raw counts to compensate for differences in sequencing depth.')
	adata = cf.adjust_raw(adata, **params)

	print('Computing initial PCA, KNN, and UMAP graph...')
	sc.pp.pca(adata, svd_solver='arpack', random_state=0)
	sc.pp.neighbors(adata, random_state=0, n_neighbors=100, method='umap')
	sc.tl.umap(adata, random_state=0)

	print('Genetic Parameters: Average log2FC greater than log2({}) from at least {} significant genes at B-H FDR of {}.'.format(2**params['logFC_threshold'], params['min_sig_genes'], params['pval']))
	# Find initial community detection resolution
	init_res = cf.find_resolution(adata, **params)
	# Get initial communities
	adata, to_cluser_ = cf.get_clusters(adata, res = init_res, initial = True, **params)
	# Recursively subcluster
	adata = cf.subcluster(adata, to_cluser_, **params)
	# Ensure all communities are genetically unique with respect to all other communities.
	adata = cf.do_final_de_eval(adata, **params)

	os.chdir('{}/{}'.format(path, output_folder)) # Change directory
	adata.obs['CellFindPy'].to_csv('CellFindPy_Groups.csv') # Save labels
	for ii in range(3): # Create 3 levels of communities for better interpretation.
		adata.obs['CellFindPy_Level {}'.format(ii+1)] = ['.'.join(s.split('.')[:int(ii+1)]) for s in adata.obs['CellFindPy'].tolist()]
		adata.obs['CellFindPy_Level {}'.format(ii+1)].to_csv('CellFindPy_Level_{}.csv'.format(ii+1)) # Save labels

	sc.settings.set_figure_params(dpi=100, dpi_save=300, format='png', frameon=False, transparent=True, fontsize=16)
	for obs in list(adata.obs.keys()): # Make umap of every obs value
		sc.pl.umap(adata, color=obs, save='{}.png'.format(obs), show=False)

	if params['sheets'] == True:
		cf.generate_gene_output(adata, **params) # Generate DE excel sheet (Time intensive)

if __name__ == '__main__':
	parser = argparse.ArgumentParser(description = 'CellFindPy is a recursive community detection package \
	purpose built for biologist analyzing single cell RNA sequencing data. It outperforms leiden and louvain algorithms by \
	implementing recursive community detection based of biologically relevant parameters and does away with tuning the resolution parameter. \
	To use CellFindPy please cd into the directory where your AnnData file is located and call CellFindPy from the command line.')

	parser.add_argument('-i', '--AnnData', help = "Provide Anndata file name.")
	parser.add_argument('-o', '--output_folder', default = '~/', help ="Provide output folder.")
	parser.add_argument("-s", "--sheets", nargs="?", default=True, help="True or False: Do you want to produce gene expression excel sheets?")
	parser.add_argument("-c", "--community", nargs="?", default="leiden", help="leiden or louvain community detection? (leiden recommended)")
	parser.add_argument("-b", "--batch", nargs="?", default=False, help="True or False: Do you want rank genes to adjust for batch?")
	parser.add_argument("-bk", "--batch_key", nargs="?", default='', help="If you want rank genes to adjust for batch, what is the batch key in Anndata.obs?")
	parser.add_argument("-mcnts", "--min_counts", nargs="?", default=5, help="Genes with less counts than this threshold will be filtered out.")
	parser.add_argument("-mcells", "--min_cells", nargs="?", default=5, help="Genes found in less than this number of cells will be filtered out.")
	parser.add_argument("-nboot", "--nbootstraps", nargs="?", default=10, help="For community stability (resampling with replacmenet), how many bootstraps?")
	parser.add_argument("-nbf", "--nbootfraction", nargs="?", default=0.8, help="For community stability (resampling with replacmenet), Minimum fraction of cells consistently clustering together?")
	parser.add_argument("-lg2FC", "--log2FC_threshold", nargs="?", default=np.log2(1.25), help="log2(value) threshold for average expression difference of all significant genes.")
	parser.add_argument("-p", "--pval", nargs="?", default=10**-6, help="False Discovery Rate.")
	parser.add_argument("-msg", "--min_sig_genes", nargs="?", default=20, help="Minimum number of significant for each community.")
	parser.add_argument("-ncom", "--n_communities", nargs="?", default=2, help="Initial number of communities to opitmize resolution to?")
	parser.add_argument("-cfrac", "--commfraction", nargs="?", default=0.6, help="Maximum size of largest community during resolution optimization?")
	parser.add_argument("-ng", "--ngenes", nargs="?", default=20, help="Number of genes from each community to make UMAP projections for?")
	args = parser.parse_args()

	path = os.getcwd() # Get current directory
	# Run main function
	main(AnnData=args.AnnData.split(' ')[-1], output_folder=args.output_folder.split(' ')[-1],
	sheets=args.sheets, community=args.community, batch=args.batch, batch_key=args.batch_key,
	min_counts=args.min_counts, min_cells=args.min_cells, bootstraps=args.nbootstraps, threshold_cells=args.nbootfraction,
	logFC_threshold=args.log2FC_threshold, pval=args.pval, min_sig_genes=args.min_sig_genes, n_genes=args.ngenes,
	ncom=args.n_communities, cfrac=args.commfraction)

	print('CellFindPy Complete.')
	print('All output files are located at {}/{}/'.format(path, args.output_folder.split(' ')[-1]))
	os.chdir(path) # Change directory back


import os
import cv2
import sys
import argparse

from utils.coco_utils import COCO_test_helper
import numpy as np

# add path
realpath = os.path.abspath(__file__)
print(realpath)
_sep = os.path.sep
realpath = realpath.split(_sep)
sys.path.append(os.path.join(realpath[0]+_sep, *realpath[1:-1]))
sys.path.append(os.path.join(realpath[0]+_sep, *realpath[1:realpath.index('rknn_model_zoo')+1]))

OBJ_THRESH = 0.1
NMS_THRESH = 0.65
IMG_SIZE = (640, 640)  # (width, height), such as (1280, 736)

CLASSES = ("person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","truck ","boat","traffic light",
           "fire hydrant","stop sign ","parking meter","bench","bird","cat","dog ","horse ","sheep","cow","elephant",
           "bear","zebra ","giraffe","backpack","umbrella","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite",
           "baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle","wine glass","cup","fork","knife ",
           "spoon","bowl","banana","apple","sandwich","orange","broccoli","carrot","hot dog","pizza ","donut","cake","chair","sofa",
           "pottedplant","bed","diningtable","toilet ","tvmonitor","laptop	","mouse	","remote ","keyboard ","cell phone","microwave ",
           "oven ","toaster","sink","refrigerator ","book","clock","vase","scissors ","teddy bear ","hair drier", "toothbrush ")

coco_id_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
         35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
         64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]

CLASSESCOCO = ('__background__', 'person', 'bicycle', 'car', 'motorbike', 'aeroplane', 'bus', 'train', 'truck', 'boat',
           'traffic light', 'fire hydrant', '???', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
           'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', '???', 'backpack', 'umbrella', '???', '???',
           'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
           'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', '???', 'wine glass', 'cup', 'fork',
           'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
           'donut', 'cake', 'chair', 'sofa', 'pottedplant', 'bed', '???', 'diningtable', '???', '???', 'toilet',
           '???', 'tvmonitor', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
           'refrigerator', '???', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush')

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def xywh2xyxy(x):
    # Convert [x, y, w, h] to [x1, y1, x2, y2]
    y = np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y

def process(input, mask, anchors):

    anchors = [anchors[i] for i in mask]
    grid_h, grid_w = map(int, input.shape[0:2])

    box_confidence = sigmoid(input[..., 4])
    box_confidence = np.expand_dims(box_confidence, axis=-1)

    box_class_probs = sigmoid(input[..., 5:])

    box_xy = sigmoid(input[..., :2])*2 - 0.5

    col = np.tile(np.arange(0, grid_w), grid_h).reshape(-1, grid_w)
    row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_w)
    col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    row = row.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
    grid = np.concatenate((col, row), axis=-1)
    box_xy += grid
    box_xy *= (int(IMG_SIZE[1]/grid_h), int(IMG_SIZE[0]/grid_w))

    box_wh = pow(sigmoid(input[..., 2:4])*2, 2)
    box_wh = box_wh * anchors

    box = np.concatenate((box_xy, box_wh), axis=-1)

    return box, box_confidence, box_class_probs

def filter_boxes(boxes, box_confidences, box_class_probs):
    """Filter boxes with object threshold.

    # Arguments
        boxes: ndarray, boxes of objects.
        box_confidences: ndarray, confidences of objects.
        box_class_probs: ndarray, class_probs of objects.

    # Returns
        boxes: ndarray, filtered boxes.
        classes: ndarray, classes for boxes.
        scores: ndarray, scores for boxes.
    """
    box_scores = box_confidences * box_class_probs
    box_classes = np.argmax(box_scores, axis=-1)
    box_class_scores = np.max(box_scores, axis=-1)
    pos = np.where(box_class_scores >= OBJ_THRESH)

    boxes = boxes[pos]
    classes = box_classes[pos]
    scores = box_class_scores[pos]

    return boxes, classes, scores

def nms_boxes(boxes, scores):
    """Suppress non-maximal boxes.

    # Arguments
        boxes: ndarray, boxes of objects.
        scores: ndarray, scores of objects.

    # Returns
        keep: ndarray, index of effective boxes.
    """
    x = boxes[:, 0]
    y = boxes[:, 1]
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]

    areas = w * h
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        xx1 = np.maximum(x[i], x[order[1:]])
        yy1 = np.maximum(y[i], y[order[1:]])
        xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
        yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])

        w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
        h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
        inter = w1 * h1

        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= NMS_THRESH)[0]
        order = order[inds + 1]
    keep = np.array(keep)
    return keep

def yolov5_post_process(input_data):
    masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
    anchors = [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45],
              [59, 119], [116, 90], [156, 198], [373, 326]]

    boxes, classes, scores = [], [], []
    for input,mask in zip(input_data, masks):
        b, c, s = process(input, mask, anchors)
        b, c, s = filter_boxes(b, c, s)
        boxes.append(b)
        classes.append(c)
        scores.append(s)

    boxes = np.concatenate(boxes)
    boxes = xywh2xyxy(boxes)
    classes = np.concatenate(classes)
    scores = np.concatenate(scores)

    nboxes, nclasses, nscores = [], [], []
    for c in set(classes):
        inds = np.where(classes == c)
        b = boxes[inds]
        c = classes[inds]
        s = scores[inds]

        keep = nms_boxes(b, s)

        nboxes.append(b[keep])
        nclasses.append(c[keep])
        nscores.append(s[keep])

    if not nclasses and not nscores:
        return None, None, None

    boxes = np.concatenate(nboxes)
    classes = np.concatenate(nclasses)
    scores = np.concatenate(nscores)

    return boxes, classes, scores

def draw(image, boxes, scores, classes):
    """Draw the boxes on the image.

    # Argument:
        image: original image.
        boxes: ndarray, boxes of objects.
        classes: ndarray, classes of objects.
        scores: ndarray, scores of objects.
        all_classes: all classes name.
    """
    for box, score, cl in zip(boxes, scores, classes):
        top, left, right, bottom = box
        print('class: {}, score: {}'.format(CLASSES[cl], score))
        print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom))
        top = int(top)
        left = int(left)
        right = int(right)
        bottom = int(bottom)

        cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
        cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
                    (top, left - 6),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6, (0, 0, 255), 2)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Process some integers.')
    # basic params
    parser.add_argument('--model_path', type=str, required= True, help='model path, could be .pt or .rknn file')
    parser.add_argument('--img_show', action='store_true', default=False, help='draw the result and show')
    parser.add_argument('--target', type=str, default='rk1808', help='target platform')
    parser.add_argument('--device_id', type=str, default=None, help='device id')

    # data params
    parser.add_argument('--anno_json', type=str, default='../../../../../datasets/COCO/annotations/instances_val2017.json', help='coco annotation path')
    parser.add_argument('--img_folder', type=str, default='../../../../../datasets/COCO//val2017', help='img folder path')
    parser.add_argument('--coco_map_test', action='store_true', help='enable coco map test')
    args = parser.parse_args()

    # init model
    model_path = args.model_path
    if model_path.endswith('.pt') or model_path.endswith('.torchscript'):
        platform = 'pytorch'
        from common.framework_excuter.pytorch_excute import Torch_model_container
        model = Torch_model_container(args.model_path)
    elif model_path.endswith('.rknn'):
        platform = 'rknn'
        from common.framework_excuter.rknn_excute import RKNN_model_container
        model = RKNN_model_container(args.model_path, args.target, args.device_id)
    else:
        assert False, "{} is not rknn/pytorch model".format(model_path)
    print('Model-{} is {} model, starting val'.format(model_path, platform))

    img_list = os.listdir(args.img_folder)
    co_helper = COCO_test_helper(enable_letter_box=True)

    # run test
    for i in range(len(img_list)):
        print('finish {}/{}'.format(i+1, len(img_list)), end='\r')

        img_name = img_list[i]
        img = cv2.imread(os.path.join(args.img_folder, img_name))
        img = co_helper.letter_box(im= img.copy(), new_shape=(IMG_SIZE[1], IMG_SIZE[0]))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # preprocee if not rknn model
        if platform == 'pytorch':
            input_data = img.transpose((2,0,1))
            input_data = input_data.reshape(1,*input_data.shape).astype(np.float32)
            input_data = input_data/255
        else:
            input_data = img

        outputs = model.run([input_data])

        # proprocess result
        outputs = [output.reshape([3,-1]+list(output.shape[-2:])) for output in outputs]
        outputs = [np.transpose(output, (2,3,0,1)) for output in outputs]

        boxes, classes, scores = yolov5_post_process(outputs)

        if args.img_show is True:
            img_p = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            if boxes is not None:
                draw(img_p, boxes, scores, classes)
            cv2.imshow("full post process result", img_p)
            cv2.waitKeyEx(0)

        if args.coco_map_test is True:
            if boxes is not None:
                for i in range(boxes.shape[0]):
                    co_helper.add_single_record(image_id = int(img_name.split('.')[0]),
                                                category_id = coco_id_list[int(classes[i])],
                                                bbox = boxes[i],
                                                score = round(scores[i], 5).astype(np.float)
                                                )
    if args.coco_map_test is True:
        pred_json = model_path.split('.')[0]+ '_{}'.format(platform) +'.json'
        co_helper.export_to_json(pred_json)

        from utils.coco_utils import coco_eval_with_json
        coco_eval_with_json(args.anno_json, pred_json)


# -*- coding: utf-8 -*-

from cytoolz import sliding_window
from itertools import chain, permutations
from collections import namedtuple
import logging
import os

from natto import MeCab
import networkx
import matplotlib.pyplot as plt

class Sentence:
    def __init__(self):
        # mecab setup
        os.environ['MECAB_PATH'] = '/usr/local/lib/libmecab.dylib'
        os.environ['MECAB_CHARSET'] = 'utf-8'

        self.Node = namedtuple('Node', ('posid', 'surface'))
        self.START_NODE = self.Node(posid=-1, surface='')
        self.END_NODE = self.Node(posid=-2, surface='')

    def learn(self, sentences):
        logging.debug("learn")
        posid_list = self._get_posid_seq(sentences)

        # create grammer network.
        self._grammer_graph = self._create_grammer_network(posid_list)
        # self._draw()

    def generate(self, similar_words):
        logging.debug("generate{0}".format(similar_words))
        posids = self._get_posid(similar_words)
        construnts = []
        for posid in posids:
            if posid in self._grammer_graph:
                construnts.append(posid)

        # posid!,0で文字列を生成、最短文字列を採用する。
        pos_via_points = permutations(construnts)
        sentences = []
        for pos_via_point in pos_via_points:
            # costが低い方が発生率が高い
            cost, path = self._search_grammer_path(pos_via_point)
            sentence = ''
            for node in path:
                sentence += node.surface
            logging.debug("---------found paths---------")
            logging.debug("path:{0}".format(path))
            logging.debug("cost:{0}".format(cost))
            logging.debug("sentence:{0}".format(sentence))
            logging.debug("---------found paths---------")
            sentences.append((cost, sentence))

        # コスト最小のものを選択
        generated = min(sentences, key=(lambda x: x[0]))
        return generated[1]

    def _search_grammer_path(self, pos_via_point):
        pos_via_and_end = (self.START_NODE,) + pos_via_point + (self.END_NODE,)
        paths = []
        cost = 0
        for network_start_end in sliding_window(2, pos_via_and_end):
            path = networkx.bidirectional_dijkstra(self._grammer_graph, network_start_end[0], network_start_end[1])
            cost += path[0]
            node_path = path[1][1:]
            paths += node_path
        paths.pop()
        return cost, paths

    def _create_grammer_network(self, posid_list):
        # cout node and edge
        node = set()
        for posid in chain.from_iterable(posid_list):
            node.add(posid)

        edge_weight_dict = {}
        for sentence in posid_list:
            for one_edge in sliding_window(2, sentence):
                edge_weight_dict[one_edge] = edge_weight_dict.get(one_edge, 1) + 1

        max_weight = max(edge_weight_dict.values())
        # NetworkX上、weightはcost扱いなので、出現頻度が高いものが低コストになるようにする

        # create direct network
        graph = networkx.DiGraph()
        graph.add_nodes_from(node)
        for edge, weight in edge_weight_dict.items():
            # 最低コストを1とする
            cost = max_weight - weight + 1
            graph.add_edge(edge[0], edge[1], weight=cost)
        return graph

    def _get_posid_seq(self, sentences):
        option = '-d /usr/local/lib/mecab/dic/mecab-ipadic-neologd'
        posid_list = []
        with MeCab(options=option) as me_cab:
            for sentence in sentences:
                words = me_cab.parse(sentence, as_nodes=True)
                # 先頭が-1,終端が-2とする
                one_work_posid = [self.START_NODE]
                for word in words:
                    node = self.Node(posid=word.posid, surface=word.surface)
                    one_work_posid.append(node)

                one_work_posid.append(self.END_NODE)
                posid_list.append(one_work_posid)
        return posid_list

    def _get_posid(self, words):
        node_list = []
        option = '-d /usr/local/lib/mecab/dic/mecab-ipadic-neologd'
        with MeCab(options=option) as me_cab:
            for word in words:
                sw = me_cab.parse(word, as_nodes=True)
                for w in sw:
                    if w.posid != 0:
                        node = self.Node(posid=w.posid, surface=w.surface)
                        node_list.append(node)

        return node_list

    def _draw(self):
        networkx.draw_spring(self._grammer_graph, iterations=20, with_labels=True)
        plt.show()


#!/usr/bin/env python

import os,glob,string,sys
from numpy.distutils.core import setup, Extension
from numpy.distutils import fcompiler
from distutils.dep_util import newer

## -------- set these
KM = 28
JM = 1
IM = 1
NC_INC = '/usr/include'
NC_LIB = '/usr/lib'
##----------------------

if '--lite' in sys.argv:
    sys.argv.pop(sys.argv.index('--lite'))
    Lite = True
else:
    Lite = False

Extensions = [
    {'name':'grid',
     'dir':'src/grid'},
    #{'name':'timestep',
    # 'dir':'src/timestep'},
    {'name':'thermodyn',
     'dir':'src/thermodyn'},
    {'name':'emanuel_convection',
     'dir':'src/convection/emanuel'},
    {'name':'hard_adjustment',
     'dir':'src/convection/hard'},
    {'name':'sbm_convection',
     'dir':'src/convection/sbm'},
     {'name':'two_column_dynamics',
     'dir':'src/dynamics/two_column'},
    {'name':'axisymmetric_dynamics',
     'dir':'src/dynamics/axisymmetric'},
    {'name':'slab_ocean',
     'dir':'src/ocean/slab_ocean'},
    {'name':'ccm3_radiation',
     'dir':'src/radiation/ccm3',
     'cppflags':'-DSUN -DPLON=%i -DPLEV=%i -DPLEVR=%i' % (IM,KM,KM)},
    {'name':'cam3_radiation',
     'dir':'src/radiation/cam3',
     'cppflags':'-DPLEV=%i' % KM,
     'lib':['netcdf','netcdff'],
     'libdir': [NC_LIB],
     'incdir': [NC_INC]},
    {'name':'chou_radiation',
     'dir':'src/radiation/chou'},
    {'name':'greygas_radiation',
     'dir':'src/radiation/greygas'},
    {'name':'newgreygas_radiation',
     'dir':'src/radiation/greynew'},
    {'name':'ozone',
     'dir':'src/radiation/ozone'},
    {'name':'insolation',
     'dir':'src/radiation/insolation'},
    {'name':'ccm3_turbulence',
     'dir':'src/turbulence/ccm3',
     'cppflags':'-DPLON=%i -DPLEV=%i' % (IM,KM)},
    {'name':'simple_turbulence',
     'dir':'src/turbulence/simple'},
#    {'name':'rrtm_radiation_fortran',
#     'dir':'src/radiation/rrtm'},
    {'name':'gfs_dynamics',
     'dir':'src/dynamics/gfs/gfs'},
    {'name':'held_suarez',
     'dir':'src/idealised/'},
    {'name':'simple_physics',
     'dir':'src/idealised/simple_physics'},
    {'name':'dcmip',
     'dir':'src/idealised/dcmip'
    },
    {'name':'emanuel_convection_new',
     'dir':'src/convection/emanuel'}
    ]

# define extensions that will be built when the --lite option is used
LiteExtensionsNames = ['grid','timestep','insolation','ozone','thermodyn','ccm3_radiation']
ExtensionsLite = []
for ext in Extensions:
    if ext['name'] in LiteExtensionsNames:
        ExtensionsLite.append(ext)

# figure out which compiler we're goint to use
compiler = fcompiler.get_default_fcompiler()
for i in range(len(sys.argv)):
    if '--fcompiler' in sys.argv[i]:
        compiler = sys.argv.pop(i)
        compiler = compiler[compiler.index('=')+1:]
print 'Using %s compiler' % compiler

# set some fortran compiler-dependent flags
if compiler == 'gnu95':
    f77flags='-ffixed-line-length-132 -fdefault-real-8'
    f90flags='-fdefault-real-8'
elif compiler == 'intel' or compiler == 'intelem':
    f77flags='-132 -r8 -w95 -w90 -mp'
    f90flags='-r8 -w95 -mp'
elif compiler == 'ibm':
    f77flags='-qautodbl=dbl4 -qsuffix=f=f:cpp=F -qfixed=132'
    f90flags='-qautodbl=dbl4 -qsuffix=f=f90:cpp=F90 -qfree=f90'
else:
    print 'Sorry, compiler %s not supported' % compiler

for ExtList in [Extensions,ExtensionsLite]:
    for i in range(len(ExtList)):
        Defaults = {'cppflags':'-DIM=%i -DJM=%i -DKM=%i' % (IM,JM,KM),
                    'f77flags':f77flags,
                    'f90flags':f90flags}
        Defaults.update(ExtList[i])
        ExtList[i] = Defaults
    if compiler == 'ibm':
        for ext in ExtList:
            ext['cppflags']='-WF,'+string.join(ext['cppflags'].split(),',')

def getSources(dir):
    #Gets list of source files for extensions
    SrcFile = os.path.join(dir,'sources_in_order_of_compilation')
    if os.path.exists(SrcFile):
        Sources = open(SrcFile).readlines()
        Sources = [os.path.join(dir,s[:-1]) for s in Sources]
    else:
        Sources = []
        w = os.walk(dir)
        for ww in w:
            if 'ignore' not in ww[0]:
                for pattern in ['*.f','*.F','*.f90','*.F90']:
                    Sources += glob.glob(os.path.join(ww[0],pattern))

    return Sources

def buildNeeded(target,src):
    #Checks if source code is newer than extension, so extension needs to be rebuilt
    target = os.path.join('lib/climt',target)
    if not os.path.exists(target):
        return True
    for file in src:
        if newer(file,target):
            return True
    print 'Extension %s is up to date' % os.path.basename(target)
    return False

def build_ext(name=None, dir=None, cppflags='', f77flags='', f90flags='', \
              lib='', libdir='', incdir=''):
    #Builds an extension
    src = getSources(dir)
    target = '_%s.so' % name
    driver = glob.glob(os.path.join(dir,'Driver.f*'))[0]
    f77flags = '"%s %s"' % (cppflags,f77flags)
    f90flags = '"%s %s"' % (cppflags,f90flags)
    if buildNeeded(target,src):
        print '\n Building %s ... \n' % os.path.basename(target)
        # generate signature file
        os.system('f2py --overwrite-signature %s -m _%s -h _%s.pyf'%(driver,name,name))
        # compile extension
        F2pyCommand = []
        F2pyCommand.append('f2py -c -m _%s' % name)
        F2pyCommand.append('--fcompiler=%s' % compiler)
        F2pyCommand.append('-I%s' % dir)
        F2pyCommand.append('-I%s' % os.path.join(dir,'include'))
        F2pyCommand.append('-I%s' % os.path.join(dir,'src'))
        F2pyCommand.append('-I%s' % os.path.join(dir,'src','include'))
        if incdir is not '':
            for i in incdir:
                F2pyCommand.append('-I%s' % i)
        if libdir is not '':
            for i in libdir:
                F2pyCommand.append('-L%s' % i)
        if lib is not '':
            for i in lib:
                F2pyCommand.append('-l%s' % i)
        F2pyCommand.append('--f77flags=%s' % f77flags)
        F2pyCommand.append('--f90flags=%s' % f90flags)
        F2pyCommand.append('_%s.pyf' % name)
        F2pyCommand.append('%s' % string.join(src))
        F2pyCommand = string.join(F2pyCommand)
        print F2pyCommand
        if os.system(F2pyCommand) > 0:
            print '+++ Compilation failed'
            sys.exit()
        os.system('mv -f _%s.so lib/climt' % name)
        os.system('rm -f _%s.pyf' % name)

#For now, just try to get the dycore to build, never mind uniformity
def build_dycore(name=None, dir=None, cppflags='', f77flags='', f90flags='', \
              lib='', libdir='', incdir=''):

    print dir
    os.chdir(dir)

    os.system('make')
    os.system('cp -f _%s.so ../../../../lib/climt/' % name)
    os.chdir('../../../..')

def setupClimt():
    # Build all extensions
    for ext in Extensions:
        if ext['name'] == 'gfs_dynamics' :
            build_dycore(**ext)

        elif ext['name'] == 'held_suarez':
                os.system('cp %s/held_suarez.py lib/climt' % ext['dir'])

        elif ext['name'] == 'dcmip':
                curr_folder = os.getcwd()
                os.chdir(ext['dir'])
                os.system('python setup.py build_ext --inplace')
                os.chdir(curr_folder)
                os.system('cp %s/dcmip.so lib/climt' % ext['dir'])

        elif ext['name'] == 'simple_physics':
                curr_folder = os.getcwd()
                os.chdir(ext['dir'])
                os.system('rm _simple_physics.c _simple_physics_custom.c')
                os.system('python setup.py build_ext --inplace')
                os.chdir(curr_folder)
                os.system('cp %s/_simple_physics.so %s/simple_physics.py \
                          %s/_simple_physics_custom.so %s/simple_physics_custom.py lib/climt'\
                          %(ext['dir'],ext['dir'],ext['dir'],ext['dir']))

        elif ext['name'] == 'newgreygas_radiation':
                curr_folder = os.getcwd()
                os.chdir(ext['dir'])
                os.system('rm _grey_gas.c')
                os.system('python setup.py build_ext --inplace')
                os.chdir(curr_folder)
                os.system('cp %s/_new_grey.so lib/climt' %(ext['dir']))

        elif ext['name'] == 'emanuel_convection_new':
                curr_folder = os.getcwd()
                os.chdir(ext['dir'])
                os.system('rm _emanuel_convection.c')
                os.system('python setup.py build_ext --inplace')
                os.chdir(curr_folder)
                os.system('cp %s/_emanuel_convection_new.so lib/climt' %(ext['dir']))

        else:
            build_ext(**ext)

    # Finish the setup
    # note: setup() cannot copy directories, and falls over
    # trying to copy the CVS directory in climt/lib/data
    # workaround: make data list which specifically excludes CVS
    os.chdir('lib/climt')
    DataFiles = []
    for File in glob.glob('data/*/*'):
        if 'CVS' not in File:
            DataFiles.append(File)
    print DataFiles
    os.chdir('../..')

    setup(name         = "CliMT",
          version      = open('Version').read()[:-1],
          description  = "Climate modelling and diagnostics toolkit",
          author       = "Rodrigo Caballero",
          author_email = "rodrigo@misu.su.se",
          url          = "http://people.su.se/~rcaba/climt",
          packages    = ['climt'],
          package_dir = {'':'lib'},
          package_data = {'climt':['*.so']+DataFiles})

def setupClimtLite():
    # Build all extensions
    for ext in ExtensionsLite: build_ext(**ext)
    os.system('mkdir -p lib/climt_lite')
    ClimtLiteFiles = ['__init__.py', '__version__.py', '_ccm3_radiation.so', '_grid.so',
     '_insolation.so', '_ozone.so', '_thermodyn.so', '_timestep.so', 'component.py',
     'grid.py', 'insolation.py', 'mathutil.py', 'ozone.py', 'parameters.py',
     'radiation.py', 'state.py', 'thermodyn.py', 'utils.py', 'io.py', 'plot.py']
    for file in ClimtLiteFiles:
        os.system('cp lib/climt/%s lib/climt_lite/%s' % (file,file))
    setup(name         = "CliMT-lite",
          version      = open('Version').read()[:-1],
          description  = "Climate modelling and diagnostics toolkit, lite version",
          author       = "Rodrigo Caballero",
          author_email = "rodrigo@misu.su.se",
          url          = "http://people.su.se/~rcaba/climt",
          packages    = ['climt_lite'],
          package_dir = {'climt_lite':'lib/climt_lite'},
          package_data = {'climt_lite':['*.so']})

if Lite:
    setupClimtLite()
else:
    setupClimt()


import time
import numpy as np
import pyaudio
import config

#--------------------------------------------#
#Refactor the original code into a OO style
#Allow to start/stop the audio input sampling and buffering
#--------------------------------------------#
class Microphone:

    def __init__(self, callback):

        #Define callback that processes the signal
        self._callback = callback

    def __del__(self):

        self.stop()

    #Start the client
    def begin(self):

        self._p = pyaudio.PyAudio()
        self._frames_per_buffer = int(config.MIC_RATE / config.FPS)
        self._stream = self._p.open(format=pyaudio.paInt16,
                        channels=1,
                        rate=config.MIC_RATE,
                        input=True,
                        frames_per_buffer=self._frames_per_buffer)
        self._overflows = 0
        self._prev_ovf_time = time.time()

    def stop(self):

        self._stream.stop_stream()
        self._stream.close()
        self._p.terminate()

    #This function is called from the main loop
    def take_samples(self):

        try:
            y = np.fromstring(self._stream.read(self._frames_per_buffer, exception_on_overflow=False), dtype=np.int16)
            y = y.astype(np.float32)
            self._stream.read(self._stream.get_read_available(), exception_on_overflow=False)
            self._callback(y)
        except IOError:
            self._overflows += 1
            if time.time() > self._prev_ovf_time + 1:
                self._prev_ovf_time = time.time()
                print('Audio buffer has overflowed {} times'.format(self._overflows))

#Original function. Todo: remove
def start_stream(callback):
    p = pyaudio.PyAudio()
    frames_per_buffer = int(config.MIC_RATE / config.FPS)
    stream = p.open(format=pyaudio.paInt16,
                    channels=1,
                    rate=config.MIC_RATE,
                    input=True,
                    frames_per_buffer=frames_per_buffer)
    overflows = 0
    prev_ovf_time = time.time()
    while True:
        try:
            y = np.fromstring(stream.read(frames_per_buffer, exception_on_overflow=False), dtype=np.int16)
            y = y.astype(np.float32)
            stream.read(stream.get_read_available(), exception_on_overflow=False)
            callback(y)
        except IOError:
            overflows += 1
            if time.time() > prev_ovf_time + 1:
                prev_ovf_time = time.time()
                print('Audio buffer has overflowed {} times'.format(overflows))
    stream.stop_stream()
    stream.close()
    p.terminate()


"""mtcnn_trt.py
"""

import numpy as np
import cv2
import pytrt

PIXEL_MEAN = 127.5
PIXEL_SCALE = 0.0078125

MINSIZE = 40

if MINSIZE == 40:
    input_h_offsets = (0, 216, 370, 478, 556, 610, 648, 676, 696)
    output_h_offsets = (0, 108, 185, 239, 278, 305, 324, 338, 348)
    max_n_scales = 9
    shape1 = (3, 710, 384)
    shape2 = (2, 350, 187)
    shape3 = (4, 350, 187)

elif MINSIZE == 120:
    input_h_offsets = (0, 72, 123, 159, 185, 203, 216, 225, 231)
    output_h_offsets = (0, 36, 61, 79, 92, 101, 108, 112, 115)
    max_n_scales = 9
    shape1 = (3, 236, 128)
    shape2 = (2, 113, 59)
    shape3 = (4, 113, 59)

elif MINSIZE == 240:
    input_h_offsets = (0, 62, 80, 93, 102, 108, 113, 116, 118)
    output_h_offsets = (0, 31, 40, 46, 51, 54, 56, 58, 59)
    max_n_scales = 9
    shape1 = (3, 119, 64)
    shape2 = (2, 55, 27)
    shape3 = (4, 55, 27)

class TrtMTCNNWrapper:

    LANDMARKS = ["left_eye", "mouth_left", "nose", "right_eye", "mouth_right"]

    def __init__(self, pnet, rnet, onet):
        self.mtcnn = TrtMtcnn(pnet, rnet, onet)

    def detect_faces(self, img):
        result = []

        boxes, landmarks = self.mtcnn.detect(img[:, :, ::-1])
        for box, features in zip(boxes, landmarks):
            face = {}

            face["confidence"] = box[-1]

            x1, y1, x2, y2 = [int(coord) for coord in box[:-1]]
            face["box"] = [x1, y1, x2 - x1, y2 - y1]

            x_key, y_key = features[:5], features[5:]
            face["keypoints"] = {
                feat: (int(y), int(x)) for feat, x, y in zip(self.LANDMARKS, x_key, y_key)
            }

            result.append(face)

        return result

def convert_to_1x1(boxes):
    """Convert detection boxes to 1:1 sizes

    # Arguments
        boxes: numpy array, shape (n,5), dtype=float32

    # Returns
        boxes_1x1
    """
    boxes_1x1 = boxes.copy()
    hh = boxes[:, 3] - boxes[:, 1] + 1.
    ww = boxes[:, 2] - boxes[:, 0] + 1.
    mm = np.maximum(hh, ww)
    boxes_1x1[:, 0] = boxes[:, 0] + ww * 0.5 - mm * 0.5
    boxes_1x1[:, 1] = boxes[:, 1] + hh * 0.5 - mm * 0.5
    boxes_1x1[:, 2] = boxes_1x1[:, 0] + mm - 1.
    boxes_1x1[:, 3] = boxes_1x1[:, 1] + mm - 1.
    boxes_1x1[:, 0:4] = np.fix(boxes_1x1[:, 0:4])
    return boxes_1x1

def crop_img_with_padding(img, box, padding=0):
    """Crop a box from image, with out-of-boundary pixels padded

    # Arguments
        img: img as a numpy array, shape (H, W, 3)
        box: numpy array, shape (5,) or (4,)
        padding: integer value for padded pixels

    # Returns
        cropped_im: cropped image as a numpy array, shape (H, W, 3)
    """
    img_h, img_w, _ = img.shape
    if box.shape[0] == 5:
        cx1, cy1, cx2, cy2, _ = box.astype(int)
    elif box.shape[0] == 4:
        cx1, cy1, cx2, cy2 = box.astype(int)
    else:
        raise ValueError
    cw = cx2 - cx1 + 1
    ch = cy2 - cy1 + 1
    cropped_im = np.zeros((ch, cw, 3), dtype=np.uint8) + padding
    ex1 = max(0, -cx1)  # ex/ey's are the destination coordinates
    ey1 = max(0, -cy1)
    ex2 = min(cw, img_w - cx1)
    ey2 = min(ch, img_h - cy1)
    fx1 = max(cx1, 0)  # fx/fy's are the source coordinates
    fy1 = max(cy1, 0)
    fx2 = min(cx2+1, img_w)
    fy2 = min(cy2+1, img_h)
    cropped_im[ey1:ey2, ex1:ex2, :] = img[fy1:fy2, fx1:fx2, :]
    return cropped_im

def nms(boxes, threshold, type='Union'):
    """Non-Maximum Supression

    # Arguments
        boxes: numpy array [:, 0:5] of [x1, y1, x2, y2, score]'s
        threshold: confidence/score threshold, e.g. 0.5
        type: 'Union' or 'Min'

    # Returns
        A list of indices indicating the result of NMS
    """
    if boxes.shape[0] == 0:
        return []
    xx1, yy1, xx2, yy2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
    areas = np.multiply(xx2-xx1+1, yy2-yy1+1)
    sorted_idx = boxes[:, 4].argsort()

    pick = []
    while len(sorted_idx) > 0:
        # In each loop, pick the last box (highest score) and remove
        # all other boxes with IoU over threshold
        tx1 = np.maximum(xx1[sorted_idx[-1]], xx1[sorted_idx[0:-1]])
        ty1 = np.maximum(yy1[sorted_idx[-1]], yy1[sorted_idx[0:-1]])
        tx2 = np.minimum(xx2[sorted_idx[-1]], xx2[sorted_idx[0:-1]])
        ty2 = np.minimum(yy2[sorted_idx[-1]], yy2[sorted_idx[0:-1]])
        tw = np.maximum(0.0, tx2 - tx1 + 1)
        th = np.maximum(0.0, ty2 - ty1 + 1)
        inter = tw * th
        if type == 'Min':
            iou = inter / \
                  np.minimum(areas[sorted_idx[-1]], areas[sorted_idx[0:-1]])
        else:
            iou = inter / \
                  (areas[sorted_idx[-1]] + areas[sorted_idx[0:-1]] - inter)
        pick.append(sorted_idx[-1])
        sorted_idx = sorted_idx[np.where(iou <= threshold)[0]]
    return pick

def generate_pnet_bboxes(conf, reg, scale, t):
    """
    # Arguments
        conf: softmax score (face or not) of each grid
        reg: regression values of x1, y1, x2, y2 coordinates.
             The values are normalized to grid width (12) and
             height (12).
        scale: scale-down factor with respect to original image
        t: confidence threshold

    # Returns
        A numpy array of bounding box coordinates and the
        cooresponding scores: [[x1, y1, x2, y2, score], ...]

    # Notes
        Top left corner coordinates of each grid is (x*2, y*2),
        or (x*2/scale, y*2/scale) in the original image.
        Bottom right corner coordinates is (x*2+12-1, y*2+12-1),
        or ((x*2+12-1)/scale, (y*2+12-1)/scale) in the original
        image.
    """
    conf = conf.T  # swap H and W dimensions
    dx1 = reg[0, :, :].T
    dy1 = reg[1, :, :].T
    dx2 = reg[2, :, :].T
    dy2 = reg[3, :, :].T
    (x, y) = np.where(conf >= t)
    if len(x) == 0:
        return np.zeros((0, 5), np.float32)

    score = np.array(conf[x, y]).reshape(-1, 1)          # Nx1
    reg = np.array([dx1[x, y], dy1[x, y],
                    dx2[x, y], dy2[x, y]]).T * 12.       # Nx4
    topleft = np.array([x, y], dtype=np.float32).T * 2.  # Nx2
    bottomright = topleft + np.array([11., 11.], dtype=np.float32)  # Nx2
    boxes = (np.concatenate((topleft, bottomright), axis=1) + reg) / scale
    boxes = np.concatenate((boxes, score), axis=1)       # Nx5
    # filter bboxes which are too small
    #boxes = boxes[boxes[:, 2]-boxes[:, 0] >= 12., :]
    #boxes = boxes[boxes[:, 3]-boxes[:, 1] >= 12., :]
    return boxes

def generate_rnet_bboxes(conf, reg, pboxes, t):
    """
    # Arguments
        conf: softmax score (face or not) of each box
        reg: regression values of x1, y1, x2, y2 coordinates.
             The values are normalized to box width and height.
        pboxes: input boxes to RNet
        t: confidence threshold

    # Returns
        boxes: a numpy array of box coordinates and cooresponding
               scores: [[x1, y1, x2, y2, score], ...]
    """
    boxes = pboxes.copy()  # make a copy
    assert boxes.shape[0] == conf.shape[0]
    boxes[:, 4] = conf  # update 'score' of all boxes
    boxes = boxes[conf >= t, :]
    reg = reg[conf >= t, :]
    ww = (boxes[:, 2]-boxes[:, 0]+1).reshape(-1, 1)  # x2 - x1 + 1
    hh = (boxes[:, 3]-boxes[:, 1]+1).reshape(-1, 1)  # y2 - y1 + 1
    boxes[:, 0:4] += np.concatenate((ww, hh, ww, hh), axis=1) * reg
    return boxes

def generate_onet_outputs(conf, reg_boxes, reg_marks, rboxes, t):
    """
    # Arguments
        conf: softmax score (face or not) of each box
        reg_boxes: regression values of x1, y1, x2, y2
                   The values are normalized to box width and height.
        reg_marks: regression values of the 5 facial landmark points
        rboxes: input boxes to ONet (already converted to 2x1)
        t: confidence threshold

    # Returns
        boxes: a numpy array of box coordinates and cooresponding
               scores: [[x1, y1, x2, y2,... , score], ...]
        landmarks: a numpy array of facial landmark coordinates:
                   [[x1, x2, ..., x5, y1, y2, ..., y5], ...]
    """
    boxes = rboxes.copy()  # make a copy
    assert boxes.shape[0] == conf.shape[0]
    boxes[:, 4] = conf
    boxes = boxes[conf >= t, :]
    reg_boxes = reg_boxes[conf >= t, :]
    reg_marks = reg_marks[conf >= t, :]
    xx = boxes[:, 0].reshape(-1, 1)
    yy = boxes[:, 1].reshape(-1, 1)
    ww = (boxes[:, 2]-boxes[:, 0]).reshape(-1, 1)
    hh = (boxes[:, 3]-boxes[:, 1]).reshape(-1, 1)
    marks = np.concatenate((xx, xx, xx, xx, xx, yy, yy, yy, yy, yy), axis=1)
    marks += np.concatenate((ww, ww, ww, ww, ww, hh, hh, hh, hh, hh), axis=1) * reg_marks
    ww = ww + 1
    hh = hh + 1
    boxes[:, 0:4] += np.concatenate((ww, hh, ww, hh), axis=1) * reg_boxes
    return boxes, marks

def clip_dets(dets, img_w, img_h):
    """Round and clip detection (x1, y1, ...) values.

    Note we exclude the last value of 'dets' in computation since
    it is 'conf'.
    """
    dets[:, 0:-1] = np.fix(dets[:, 0:-1])
    evens = np.arange(0, dets.shape[1]-1, 2)
    odds  = np.arange(1, dets.shape[1]-1, 2)
    dets[:, evens] = np.clip(dets[:, evens], 0., float(img_w-1))
    dets[:, odds]  = np.clip(dets[:, odds], 0., float(img_h-1))
    return dets

class TrtPNet(object):
    """TrtPNet

    Refer to mtcnn/det1_relu.prototxt for calculation of input/output
    dimmensions of TrtPNet, as well as input H offsets (for all scales).
    The output H offsets are merely input offsets divided by stride (2).
    """

    def __init__(self, engine):
        """__init__

        # Arguments
            engine: path to the TensorRT engine file
        """
        self.trtnet = pytrt.PyTrtMtcnn(engine, shape1, shape2, shape3)
        self.trtnet.set_batchsize(1)

    def detect(self, img, factor=0.709, threshold=0.7):
        """Detect faces using PNet

        # Arguments
            img: input image as a RGB numpy array
            threshold: confidence threshold

        # Returns
            A numpy array of bounding box coordinates and the
            cooresponding scores: [[x1, y1, x2, y2, score], ...]
        """

        if MINSIZE < 40:
            raise ValueError("TrtPNet is currently designed with "
                             "'minsize' >= 40")
        if factor > 0.709:
            raise ValueError("TrtPNet is currently designed with "
                             "'factor' <= 0.709")
        m = 12.0 / MINSIZE
        img_h, img_w, _ = img.shape
        minl = min(img_h, img_w) * m

        # create scale pyramid
        scales = []
        while minl >= 12:
            scales.append(m)
            m *= factor
            minl *= factor
        if len(scales) > max_n_scales:  # probably won't happen...
            raise ValueError('Too many scales, try increasing minsize '
                             'or decreasing factor.')

        total_boxes = np.zeros((0, 5), dtype=np.float32)
        img = (img.astype(np.float32) - PIXEL_MEAN) * PIXEL_SCALE

        # stack all scales of the input image vertically into 1 big
        # image, and only do inferencing once
        im_data = np.zeros((1, *shape1), dtype=np.float32)
        for i, scale in enumerate(scales):
            h_offset = input_h_offsets[i]
            h = int(img_h * scale)
            w = int(img_w * scale)
            im_data[0, :, h_offset:(h_offset+h), :w] = \
                cv2.resize(img, (w, h)).transpose((2, 0, 1))

        out = self.trtnet.forward(im_data)

        # extract outputs of each scale from the big output blob
        for i, scale in enumerate(scales):
            h_offset = output_h_offsets[i]
            h = (int(img_h * scale) - 12) // 2 + 1
            w = (int(img_w * scale) - 12) // 2 + 1
            pp = out['prob1'][0, 1, h_offset:(h_offset+h), :w]
            cc = out['boxes'][0, :, h_offset:(h_offset+h), :w]
            boxes = generate_pnet_bboxes(pp, cc, scale, threshold)
            if boxes.shape[0] > 0:
                pick = nms(boxes, 0.5, 'Union')
                if len(pick) > 0:
                    boxes = boxes[pick, :]
            if boxes.shape[0] > 0:
                total_boxes = np.concatenate((total_boxes, boxes), axis=0)

        if total_boxes.shape[0] == 0:
            return total_boxes
        pick = nms(total_boxes, threshold, 'Union')
        dets = clip_dets(total_boxes[pick, :], img_w, img_h)
        return dets

    def destroy(self):
        self.trtnet.destroy()
        self.trtnet = None

class TrtRNet(object):
    """TrtRNet

    # Arguments
        engine: path to the TensorRT engine (det2) file
    """

    def __init__(self, engine):
        self.trtnet = pytrt.PyTrtMtcnn(engine,
                                       (3, 24, 24),
                                       (2, 1, 1),
                                       (4, 1, 1))

    def detect(self, img, boxes, max_batch=32, threshold=0.7):
        """Detect faces using RNet

        # Arguments
            img: input image as a RGB numpy array
            boxes: detection results by PNet, a numpy array [:, 0:5]
                   of [x1, y1, x2, y2, score]'s
            max_batch: only process these many top boxes from PNet
            threshold: confidence threshold

        # Returns
            A numpy array of bounding box coordinates and the
            cooresponding scores: [[x1, y1, x2, y2, score], ...]
        """
        if max_batch > 256:
            raise ValueError('Bad max_batch: %d' % max_batch)
        boxes = boxes[:max_batch]  # assuming boxes are sorted by score
        if boxes.shape[0] == 0:
            return boxes
        img_h, img_w, _ = img.shape
        boxes = convert_to_1x1(boxes)
        crops = np.zeros((boxes.shape[0], 24, 24, 3), dtype=np.uint8)
        for i, det in enumerate(boxes):
            cropped_im = crop_img_with_padding(img, det)
            # NOTE: H and W dimensions need to be transposed for RNet!
            crops[i, ...] = cv2.transpose(cv2.resize(cropped_im, (24, 24)))
        crops = crops.transpose((0, 3, 1, 2))  # NHWC -> NCHW
        crops = (crops.astype(np.float32) - PIXEL_MEAN) * PIXEL_SCALE

        self.trtnet.set_batchsize(crops.shape[0])
        out = self.trtnet.forward(crops)

        pp = out['prob1'][:, 1, 0, 0]
        cc = out['boxes'][:, :, 0, 0]
        boxes = generate_rnet_bboxes(pp, cc, boxes, threshold)
        if boxes.shape[0] == 0:
            return boxes
        pick = nms(boxes, 0.7, 'Union')
        dets = clip_dets(boxes[pick, :], img_w, img_h)
        return dets

    def destroy(self):
        self.trtnet.destroy()
        self.trtnet = None

class TrtONet(object):
    """TrtONet

    # Arguments
        engine: path to the TensorRT engine (det3) file
    """

    def __init__(self, engine):
        self.trtnet = pytrt.PyTrtMtcnn(engine,
                                       (3, 48, 48),
                                       (2, 1, 1),
                                       (4, 1, 1),
                                       (10, 1, 1))

    def detect(self, img, boxes, max_batch=64, threshold=0.7):
        """Detect faces using ONet

        # Arguments
            img: input image as a RGB numpy array
            boxes: detection results by RNet, a numpy array [:, 0:5]
                   of [x1, y1, x2, y2, score]'s
            max_batch: only process these many top boxes from RNet
            threshold: confidence threshold

        # Returns
            dets: boxes and conf scores
            landmarks
        """
        if max_batch > 64:
            raise ValueError('Bad max_batch: %d' % max_batch)
        if boxes.shape[0] == 0:
            return (np.zeros((0, 5), dtype=np.float32),
                    np.zeros((0, 10), dtype=np.float32))
        boxes = boxes[:max_batch]  # assuming boxes are sorted by score
        img_h, img_w, _ = img.shape
        boxes = convert_to_1x1(boxes)
        crops = np.zeros((boxes.shape[0], 48, 48, 3), dtype=np.uint8)
        for i, det in enumerate(boxes):
            cropped_im = crop_img_with_padding(img, det)
            # NOTE: H and W dimensions need to be transposed for RNet!
            crops[i, ...] = cv2.transpose(cv2.resize(cropped_im, (48, 48)))
        crops = crops.transpose((0, 3, 1, 2))  # NHWC -> NCHW
        crops = (crops.astype(np.float32) - PIXEL_MEAN) * PIXEL_SCALE

        self.trtnet.set_batchsize(crops.shape[0])
        out = self.trtnet.forward(crops)

        pp = out['prob1'][:, 1, 0, 0]
        cc = out['boxes'][:, :, 0, 0]
        mm = out['landmarks'][:, :, 0, 0]
        boxes, landmarks = generate_onet_outputs(pp, cc, mm, boxes, threshold)
        pick = nms(boxes, 0.7, 'Min')
        return (clip_dets(boxes[pick, :], img_w, img_h),
                np.fix(landmarks[pick, :]))

    def destroy(self):
        self.trtnet.destroy()
        self.trtnet = None

class TrtMtcnn(object):
    """TrtMtcnn"""

    def __init__(self, pnet, rnet, onet):
        self.pnet = TrtPNet(pnet)
        self.rnet = TrtRNet(rnet)
        self.onet = TrtONet(onet)

    def __del__(self):
        self.onet.destroy()
        self.rnet.destroy()
        self.pnet.destroy()

    def _detect_1280x720(self, img):
        """_detec_1280x720()

        Assuming 'img' has been resized to less than 1280x720.
        """
        # MTCNN model was trained with 'MATLAB' image so its channel
        # order is RGB instead of BGR.
        img = img[:, :, ::-1]  # BGR -> RGB
        dets = self.pnet.detect(img)
        dets = self.rnet.detect(img, dets)
        dets, landmarks = self.onet.detect(img, dets)
        return dets, landmarks

    def detect(self, img):
        """detect()

        This function handles rescaling of the input image if it's
        larger than 1280x720.
        """
        if img is None:
            raise ValueError
        img_h, img_w, _ = img.shape
        scale = min(720. / img_h, 1280. / img_w)
        if scale < 1.0:
            new_h = int(np.ceil(img_h * scale))
            new_w = int(np.ceil(img_w * scale))
            img = cv2.resize(img, (new_w, new_h))
            minsize = max(int(np.ceil(MINSIZE * scale)), 40)
        dets, landmarks = self._detect_1280x720(img)
        if scale < 1.0:
            dets[:, :-1] = np.fix(dets[:, :-1] / scale)
            landmarks = np.fix(landmarks / scale)
        return dets, landmarks


import torch
import numpy as np
from copy import deepcopy
from typing import Dict, Union, Optional

from tianshou.policy import BasePolicy
from tianshou.data import Batch, ReplayBuffer, to_torch_as, to_numpy

class DQNPolicy(BasePolicy):
    """Implementation of Deep Q Network. arXiv:1312.5602

    Implementation of Double Q-Learning. arXiv:1509.06461

    Implementation of Dueling DQN. arXiv:1511.06581 (the dueling DQN is
    implemented in the network side, not here)

    :param torch.nn.Module model: a model following the rules in
        :class:`~tianshou.policy.BasePolicy`. (s -> logits)
    :param torch.optim.Optimizer optim: a torch.optim for optimizing the model.
    :param float discount_factor: in [0, 1].
    :param int estimation_step: greater than 1, the number of steps to look
        ahead.
    :param int target_update_freq: the target network update frequency (``0``
        if you do not use the target network).
    :param bool reward_normalization: normalize the reward to Normal(0, 1),
        defaults to ``False``.

    .. seealso::

        Please refer to :class:`~tianshou.policy.BasePolicy` for more detailed
        explanation.
    """

    def __init__(self,
                 model: torch.nn.Module,
                 optim: torch.optim.Optimizer,
                 discount_factor: float = 0.99,
                 estimation_step: int = 1,
                 target_update_freq: Optional[int] = 0,
                 reward_normalization: bool = False,
                 grad_norm_clipping=10,
                 **kwargs) -> None:
        super().__init__(**kwargs)
        self.model = model
        self.optim = optim
        self.eps = 0
        assert 0 <= discount_factor <= 1, 'discount_factor should in [0, 1]'
        self._gamma = discount_factor
        assert estimation_step > 0, 'estimation_step should greater than 0'
        self._n_step = estimation_step
        self._target = target_update_freq > 0
        self._freq = target_update_freq
        self.grad_norm_clipping = grad_norm_clipping
        self._cnt = 0
        if self._target:
            self.model_old = deepcopy(self.model)
            self.model_old.eval()
        self._rew_norm = reward_normalization

    def set_eps(self, eps: float) -> None:
        """Set the eps for epsilon-greedy exploration."""
        self.eps = eps

    def train(self, mode=True) -> torch.nn.Module:
        """Set the module in training mode, except for the target network."""
        self.training = mode
        self.model.train(mode)
        return self

    def sync_weight(self) -> None:
        """Synchronize the weight for the target network."""
        self.model_old.load_state_dict(self.model.state_dict())

    def _target_q(self, buffer: ReplayBuffer,
                  indice: np.ndarray) -> torch.Tensor:
        batch = buffer[indice]  # batch.obs_next: s_{t+n}
        if self._target:
            # target_Q = Q_old(s_, argmax(Q_new(s_, *)))
            a = self(batch, input='obs_next', eps=0).act
            with torch.no_grad():
                target_q = self(
                    batch, model='model_old', input='obs_next').logits
            target_q = target_q[np.arange(len(a)), a]
        else:
            with torch.no_grad():
                target_q = self(batch, input='obs_next').logits.max(dim=1)[0]
        return target_q

    def process_fn(self, batch: Batch, buffer: ReplayBuffer,
                   indice: np.ndarray) -> Batch:
        """Compute the n-step return for Q-learning targets. More details can
        be found at :meth:`~tianshou.policy.BasePolicy.compute_nstep_return`.
        """
        batch = self.compute_nstep_return(
            batch, buffer, indice, self._target_q,
            self._gamma, self._n_step, self._rew_norm)
        return batch

    def forward(self, batch: Batch,
                state: Optional[Union[dict, Batch, np.ndarray]] = None,
                model: str = 'model',
                input: str = 'obs',
                eps: Optional[float] = None,
                **kwargs) -> Batch:
        """Compute action over the given batch data. If you need to mask the
        action, please add a "mask" into batch.obs, for example, if we have an
        environment that has "0/1/2" three actions:
        ::

            batch == Batch(
                obs=Batch(
                    obs="original obs, with batch_size=1 for demonstration",
                    mask=np.array([[False, True, False]]),
                    # action 1 is available
                    # action 0 and 2 are unavailable
                ),
                ...
            )

        :param float eps: in [0, 1], for epsilon-greedy exploration method.

        :return: A :class:`~tianshou.data.Batch` which has 3 keys:

            * ``act`` the action.
            * ``logits`` the network's raw output.
            * ``state`` the hidden state.

        .. seealso::

            Please refer to :meth:`~tianshou.policy.BasePolicy.forward` for
            more detailed explanation.
        """
        model = getattr(self, model)
        obs = getattr(batch, input)
        obs_ = obs.obs if hasattr(obs, 'obs') else obs
        q, h = model(obs_, state=state, info=batch.info)
        act = to_numpy(q.max(dim=1)[1])
        has_mask = hasattr(obs, 'mask')
        if has_mask:
            # some of actions are masked, they cannot be selected
            q_ = to_numpy(q)
            q_[~obs.mask] = -np.inf
            act = q_.argmax(axis=1)
        # add eps to act
        if eps is None:
            eps = self.eps
        if not np.isclose(eps, 0):
            for i in range(len(q)):
                if np.random.rand() < eps:
                    q_ = np.random.rand(*q[i].shape)
                    if has_mask:
                        q_[~obs.mask[i]] = -np.inf
                    act[i] = q_.argmax()
        return Batch(logits=q, act=act, state=h)

    def learn(self, batch: Batch, **kwargs) -> Dict[str, float]:
        if self._target and self._cnt % self._freq == 0:
            self.sync_weight()
        weight = batch.pop("weight", 1.0)
        self.optim.zero_grad()
        q = self(batch, eps=0.).logits
        q = q[np.arange(len(q)), batch.act]
        r = to_torch_as(batch.returns, q).flatten()
        c = torch.nn.SmoothL1Loss(reduction = 'none')
        # c = lambda r, q: (r-q).pow(2)
        td = c(r, q)
        loss = (td * weight).mean()
        batch.weight = r - q  # prio-buffer
        loss.backward()
        if self.grad_norm_clipping:
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_norm_clipping)
        self.optim.step()
        self._cnt += 1
        return {'loss': loss.item()}