# The original code requires pylcaio release v1.0, from Agez, found at:
# https://github.com/MaximeAgez/pylcaio

import pandas as pd
import numpy as np
import os, sys
import pickle
import logging
import pylcaio

fp_results = os.path.join(os.path.curdir, 'results')
fp_output = os.path.join(os.path.curdir, 'output')

def hybrid_emission_factors(trade_only, year):
    Analysis = pylcaio.Analysis('ecoinvent3.5','exiobase3', method_double_counting='STAM')

    # For making human-readable indices; consists of all ecoinvent process metadata
    labels = Analysis.PRO_f

    # Make DataFrame for whole ecoinvent requirements matrix with human-readable indices
    Aff_labels = (labels.join(Analysis.A_ff,how='inner'))
    Aff_labels = (Aff_labels).set_index(keys=['activityName', 'geography', 'productName'], append=True)
    Aff_labels = Aff_labels.iloc[:,19:] # remove extraeneous metadata
    Aff_labels.columns = Aff_labels.index

    # Calculate life cycle impacts
    I = pd.DataFrame(np.eye(len(Analysis.A_ff)), Analysis.A_ff.index, Analysis.A_ff.columns)
    X = pd.DataFrame(np.linalg.solve(I-Analysis.A_ff, I), Analysis.A_ff.index, I.columns) # pro x pro
    D = Analysis.C_f.fillna(0).dot(Analysis.F_f).dot(X) # pro x imp

    D = D.T
    D_labels = labels.join(D, how='inner')
    D_labels.set_index(keys=['activityName', 'geography', 'productName'], append=True, inplace=True)
    D_labels = D_labels.iloc[:, 19:]
    D_labels = D_labels.T

    # Other possible impact methods
    imp_list = ['EDIP; environmental impact; global warming, GWP 100a; kg CO2-Eq', 'ReCiPe Midpoint (H); climate change; GWP100; kg CO2-Eq', 'ReCiPe Midpoint (H) V1.13; climate change; GWP100; kg CO2-Eq', 'IPCC 2013; climate change; GTP 100a; kg CO2-Eq']
    D_labels_comp = D_labels.loc[imp_list]
    D_labels = D_labels.loc['EDIP; environmental impact; global warming, GWP 100a; kg CO2-Eq']

    # the footprint of all hybridized processes from ecoinvent, with their original footprint
    fp_hybrid = os.path.join(os.path.curdir, 'data', 'Results_full_database_STAM_2011.xlsx')
    hybrid_results = pd.read_excel(fp_hybrid, 'GWP_only_hyb', usecols='A:E,G', index_col=[0,1,2,3,4])

    # ## List of ENTSO-E countries
    country_list = ['AL','AT','BA','BE','BG','CH','CZ','DE','DK','EE',
                    'ES','FI','FR','GR','HR','HU','IE','IT','LT','LU',
                    'LV','ME','MK','NL','NO','PL','PT','RO','RS','SE',
                    'SI','SK','GB','TR','GE','CY','HR','NI','UK']

    # Find all electricity processes
    set(labels.loc[labels['activityName'].str.contains('electricity')]['activityName'])

    # Isolate high voltage (production) mixes and low voltage mixes (which includes solar PV)
    # note that pumped storage hydropower in CH gets excluded as it has 0 contribution to the production mix (for whatever reason)
    market_mixes = labels.loc[labels['activityName'].str.contains('electricity, high voltage, production mix')] # production mixes, do not include solar or waste (or Swiss pumped hydro)
    waste_mixes = labels.loc[labels['activityName'].str.contains('electricity, from municipal waste incineration to generic market for electricity, medium voltage')]  # production mixes, do not include solar
    solar_mixes = labels.loc[labels['activityName'].str.contains('electricity, low voltage')]  # use low voltage mixes to get shares of solar photovoltaic

    # ### Filter for high-voltage production mixes
    # First, retrieve electricity mixes, and reduce un-needed processes
    df = Aff_labels[market_mixes.index]
    df_waste = Aff_labels.loc[waste_mixes.index]
    df_solar = Aff_labels[solar_mixes.index]

    df = df.loc[:, df.columns.isin(country_list, level=2)]

    # manual workaround for keeping CH pumped hydro after removing unneeded rows
    # (since this ecoinvent process apparently does not contribute to any electricity mixes in ecoinvent)
    ind =('69ccf019-081e-42e3-b7f5-879dc7dde86a_66c93e71-f32b-4591-901c-55395db5c132', 'electricity production, hydro, pumped storage', 'CH', 'electricity, high voltage')  # Activityid for Swiss pumped hydro
    col = df.columns.get_loc(('54ed269b-e329-469c-902e-a8991ee7d24a_66c93e71-f32b-4591-901c-55395db5c132', 'electricity, high voltage, production mix', 'CH', 'electricity, high voltage'))
    index = df.index.get_loc(ind)
    df.iloc[index, col] = 1e-7

    df = df.replace(0, np.nan).dropna(how='all', axis=0)

    # ### Filter for medium-voltage mixes ( to find waste incineration shares)
    df_waste = df_waste.loc[:, df_waste.columns.isin(country_list, level=2)]
    df_waste = df_waste.replace(to_replace=0, value=np.nan).dropna(how='all', axis=0)

    # check unique waste incineration activities
    set(df_waste.index.get_level_values('activityName'))

    # ### Filter for low-voltage mixes (to find solar photovoltaic shares)
    df_solar = df_solar.loc[:, df_solar.columns.isin(country_list, level=2)]
    df_solar = df_solar.replace(to_replace=0, value=np.nan).dropna(how='all', axis=0)

    # ### Develop correspondence between ecoinvent and ENTSO-E technology categories

    # Correspondence from ENTSO-E technology categories to keywords for searching in ecoinvent processes
    # Note that the 'fossil oil shale', 'other' and 'other renewable' do not have ecoinvent equivalents;
    # these are dealt with later
    tec_dict = {'Biomass': ['biogas', 'wood chips'],
                'Fossil Brown coal/Lignite': 'lignite',
                'Fossil Coal-derived gas': 'coal gas',
                'Fossil Gas': 'natural gas',
                'Fossil Hard coal': 'hard coal',
                'Fossil Oil': ' oil',
                'Fossil Oil shale': '--',
                'Fossil Peat': 'peat',
                'Geothermal': 'geothermal',
                'Hydro Pumped Storage': 'pumped storage',
                'Hydro Run-of-river and poundage': 'run-of-river',
                'Hydro Water Reservoir': 'reservoir',
                'Nuclear': 'nuclear',
                'Other': '---',
                'Other renewable': '----',
                'Solar': ['solar thermal', 'solar tower', 'photovoltaic'],
                'Waste': 'waste incineration',
                'Wind Offshore': 'offshore',
                'Wind Onshore': 'onshore'
    }

    # reverse correspondence from ecoinvent search keywords to ENTSO-E categories
    temp_tecdict = { 'biogas':'Biomass',
                    'wood chips':'Biomass',
                'lignite': 'Fossil Brown coal/Lignite' ,
                'coal gas': 'Fossil Coal-derived gas',
                'natural gas': 'Fossil Gas' ,
                'hard coal': 'Fossil Hard coal' ,
                ' oil': 'Fossil Oil',
                '--' : 'Fossil Oil shale',
                'peat': 'Fossil Peat' ,
                'geothermal':'Geothermal',
                'pumped storage': 'Hydro Pumped Storage',
                'run-of-river': 'Hydro Run-of-river and poundage',
                'reservoir': 'Hydro Water Reservoir',
                'nuclear' : 'Nuclear',
                '---':'Other' ,
                '----': 'Other renewable',
                'solar thermal': 'Solar',
                'solar tower': 'Solar',
                'photovoltaic': 'Solar',
                'waste incineration': 'Waste', #  'electricity, from municipal waste incineration to generic market for electricity, medium voltage',
                'offshore': 'Wind Offshore',
                'onshore':'Wind Onshore'
    }

    # flatten tec list
    tec_list = list(tec_dict.values())
    for item in tec_list:
        if type(item)==list:
            for subitem in item:
                tec_list.append(subitem)
            tec_list.remove(item)

    # Get list of solar PV processes in ecoinvent
    d = list(set(df_solar.index.get_level_values(level=1)))
    pv = [entry for entry in d if 'electricity production' in entry]  # all low-voltage production technologies

    # Fetch waste incineration process names
    waste = ['electricity, from municipal waste incineration to generic market for electricity, medium voltage']

    # Make list of electricity generation technologies (both high voltage and low-voltage, i.e., solar PV)
    ecoinvent_tecs = list(set(df.index.get_level_values(level=1))) + pv + waste

    # Quality assurance; make sure all ecoinvent technologies are covered by the keyword search, and keyword matches are correct
    # also, check number of matches for each keyword
    matches = []
    match_keys = []
    num_matches = dict((keyword,0) for keyword in tec_list)

    for tec in ecoinvent_tecs:
        for keyword in tec_list:
            if keyword in tec and tec not in matches:
#                print(f'{keyword} is in {tec}')
                matches.append(tec)
                match_keys.append(temp_tecdict[keyword])
                num_matches[keyword] += 1

    print(num_matches)

    # Find ecoinvent tecs that were not matched
    if len(ecoinvent_tecs)-len(matches) > 0:
        print('ecoinvent process(es) not matched:')
        print(set(ecoinvent_tecs).difference(set(matches)))

    # Correspondence of ENTSO-E categories and specific ecoinvent processes
    from collections import defaultdict

    ei_tec_dict = defaultdict(list) # keys as entso-e categories, values are corresponding ei
    rev_ei_tec_dict = defaultdict(list) # vice-versa; keys as ei, values are entso-e categories
    for i, j in zip(match_keys, matches):
        ei_tec_dict[i].append(j)
        rev_ei_tec_dict[j].append(i)

    # Build correspondence table for Supplementary Information
    tec_index = pd.Index(list(ei_tec_dict.keys()))
    tec_values = list(ei_tec_dict.values())
    ei_tec_table = pd.DataFrame(tec_values, index=tec_index).stack()
    ei_tec_table.index = ei_tec_table.index.droplevel(level=1)
    ei_tec_table.sort_index(axis=0, inplace=True)

    # ### Calculate shares of ecoinvent processes
    # For multi-process technology categories, e.g., solar, biomass...

    # combine high-, medium- and low-voltage electricity mixes for full tech resolution
    f = pd.concat([df, df_waste, df_solar], axis=1)

    rev_ei_tec_dict = {k: v for k, v in rev_ei_tec_dict.items() if not not v}

    temp_list = f.index.get_level_values('activityName')
    entso_list = []
    f_check = []

    # For every activity in ecoinvent,if they are related to electricity production, add to entso_list
    # If not, drop from the matrix; they are irrelevant
    for entry in temp_list:
        try:
            entso_list.append(rev_ei_tec_dict[entry])
        except:
            f_check.append(entry)

    f.drop(index=set(f_check), level=1, axis=0, inplace=True)

    # inputs to electricity mixes not included in the calculations of emission factors
    set(f_check)

    entso_list = [item for sublist in entso_list for item in sublist] #flatten list

    entso_index = pd.Index(tuple(entso_list), name='entso_e')
    f.set_index(entso_index, append=True, inplace=True)

    # ### Get shares of each ecoinvent process for each ENTSO-E technology category

    f_gpby = f.groupby(level='entso_e')

    # Get shares of each ecoinvent process in each ENTSO-E technology category;
    # for calculating weighted average for emission factor
    ei_process_shares = f.divide(f_gpby.sum(), level=4, axis=1)

    hybrid_temp = hybrid_results.reset_index(level=[1,2,3,4], drop=True)
    g_temp = ei_process_shares.reset_index(level=[1,2,3,4], drop=True)

    # ### Calculate weighted average, hybridized emissions factor for each technology

    ef = ei_process_shares.mul(hybrid_temp.iloc[:, 0], axis=0, level=0)

    ef_aggregated = ef.groupby(level='entso_e').sum()

    ef_countries = ef_aggregated.sum(axis=1, level='geography')
    ef_countries.replace(0, np.nan, inplace=True)
    ef_countries.dropna(axis=1, how='all', inplace=True)
    ef_countries.sort_index(inplace=True)

    ef_countries = ef_countries.T
    ef_countries.sort_index(axis=0, inplace=True)

    entso_fp = os.path.join(fp_output, 'entsoe', 'ENTSO_production_volumes_' + str(year) + '.csv')

    entso_e_production = pd.read_csv(entso_fp, header=0, index_col=[0])
    entso_e_production.replace(0, np.nan,inplace=True)
    entso_mask = entso_e_production.isna().sort_index()
    ef_mask = ef_countries.sort_index().isna()

    entso_mask.drop(columns=['Fossil Oil shale', 'Other', 'Other renewable', 'Marine'], inplace=True)

    entso_mask.index.rename('geography', inplace=True)
    entso_mask.columns.rename('entso_e', inplace=True)

    # Countries for which we have production data, but no hybridized emission factors
    no_ef_countries = set(entso_e_production.index) - set(ef_countries.index)
    print('Countries with production data, but no hybridized emission factors:')
    print(no_ef_countries)
    logging.warning(f'{no_ef_countries} have production data from ENTSO_E but no hybridized emission factors')

    list(set(ef_countries.columns) - set(entso_e_production.index))

    check_df = entso_mask.eq(ef_mask, axis=0)

    # Find countries missing from ecoinvent
    missing_countries = []
    for country, row in check_df.iterrows():
        if not row.any():
            missing_countries.append(country)
    # Optionally, drop them from mask
    check_df.drop(index = missing_countries, inplace=True)

    # Find the technology-country pairs that are missing hybridized emission factors
    countries_missing_ef = {}
    for tec, col in check_df.iteritems():
        temp = check_df.index[check_df[tec] == False].tolist()
        print(f'For {tec}, regionalized, hybridized emission factors are missing for {temp}. Using technological average from available regions instead.')
        temp2 = []
        for country in temp:
            if ef_mask.loc[country, tec]:
                temp2.append(country)
        countries_missing_ef[tec] = temp2#.append(country)
    countries_missing_ef

    no_ef = pd.DataFrame({k:pd.Series(v[:13]) for k, v in countries_missing_ef.items()})

    # ## Calculate LC emissions for production mix for countries with missing production data

#    trade_only = ['AL','BY','HR','LU','MD','MT','RU','TR','UA']
    low_volt = labels.loc[labels['activityName'] == 'market for electricity, low voltage']
    trade_only_mixes_lv = low_volt.loc[low_volt['geography'].isin(trade_only)].index
    lv_mix_ef = D_labels.loc[trade_only_mixes_lv]  # low voltage mixes, not used

    high_volt = labels.loc[labels['activityName'] == 'market for electricity, high voltage']
    trade_only_mixes_hv = high_volt.loc[high_volt['geography'].isin(trade_only)].index
    hv_mix_ef = D_labels.loc[trade_only_mixes_hv]

    lv_mix_ef_comp = D_labels_comp[trade_only_mixes_lv]  # low voltage mixes, not used
    hv_mix_ef_comp = D_labels_comp[trade_only_mixes_hv]

    trade_mixes_comp = lv_mix_ef_comp.join(hv_mix_ef_comp)

    labels.loc[labels['activityName'] == 'electricity, from municipal waste incineration to generic market for electricity, medium voltage']

    # ## Export results
    # Convert emission factors to g CO2-eq/kWh
    ef_countries = ef_countries * 1000
    lv_mix_ef = lv_mix_ef * 1000  # low voltage mixes, not used
    hv_mix_ef = hv_mix_ef * 1000
    trade_mixes_comp = trade_mixes_comp * 1000

    # Calculate mean technology footprints as quick check
    ef_aggregated.replace(0, np.nan).mean(axis=1)

    fp_results = os.path.join(os.path.curdir,'results')

    with pd.ExcelWriter(os.path.join(fp_results, 'hybrid_emission_factors_final_' + str(year) + '.xlsx')) as writer:
        ef_countries.to_excel(writer, sheet_name='country_emission_factors')
        no_ef.to_excel(writer, sheet_name='missing_emission_factors')
        lv_mix_ef.to_excel(writer, sheet_name='ecoinvent ef')
        hv_mix_ef.to_excel(writer, sheet_name='ecoinvent ef_hv')
        trade_mixes_comp.to_excel(writer, sheet_name='trade_mixes_from_ei')
        ei_tec_table.to_excel(writer, sheet_name='correspondence table')
        ef_aggregated.to_excel(writer, sheet_name='mean tech fps')

    # Export to .csv for use in BEV_footprints_calculations
    # append extra countries to file
    try:
        existing_tradeonly = pd.read_csv(os.path.join(fp_output, 'ecoinvent_ef_hv.csv'))['geography'].values
        add_ef_countries = set(hv_mix_ef.index.get_level_values('geography')) - set(existing_tradeonly)  # countries not already in csv file
        if add_ef_countries:
            hv_mix_ef_exp = hv_mix_ef.loc[:,:,list(add_ef_countries),:]
            hv_mix_ef_exp.to_csv(os.path.join(fp_output, 'ecoinvent_ef_hv.csv'), mode='a', header=False)  # used in BEV_footprints_calculations
    except FileNotFoundError:
        print('Previous ecoinvent file does not exist. Making file.')
        hv_mix_ef.to_csv(os.path.join(fp_output, 'ecoinvent_ef_hv.csv'))

    ef_countries.to_csv(os.path.join(fp_output, 'country_emission_factors.csv'))  # used in clean_impact_factors
    no_ef.to_csv(os.path.join(fp_output, 'missing_emission_factors.csv'))  # used in clean_impact_factors
#    lv_mix_ef.to_csv(os.path.join(fp_results, 'ecoinvent_ef_lv.csv'))

    return ef_countries, no_ef, countries_missing_ef

def clean_impact_factors(year, trade_only, calc_hybrid):
    entsoe_fp = os.path.join(fp_output, 'entsoe')
    with open(os.path.join(entsoe_fp, 'gen_final_' + str(year) + '.pkl'), 'rb') as handle:
        gen_df = pickle.load(handle)

    with open(os.path.join(entsoe_fp,'trade_final_' + str(year) + '.pkl'), 'rb') as handle:
        trade_df = pickle.load(handle)

    if calc_hybrid:
        try:
            print('\n Calculating hybridized emission factors from pyLCAIO')
            # gen_df not modified in hybrid_emission_factors, just used to determine relevant labels
            ef, missing_factors, countries_missing_ef = hybrid_emission_factors(trade_only, year)
        except Exception as e:
            print('\n Calculating from pyLCAIO failed. Importing previously calculated emission factors instead')
            print(e)
            # import ready-calculated emission factors if no access to pylcaio object
            ef = pd.read_csv(os.path.join(fp_output, 'country_emission_factors.csv'), index_col=[0])
            missing_factors = pd.read_csv(os.path.join(fp_output, 'missing_emission_factors.csv'), index_col=[0])
    else:
        print('\n Loading previously calculated emission factors')
        ef = pd.read_csv(os.path.join(fp_output, 'country_emission_factors.csv'), index_col=[0])
        missing_factors = pd.read_csv(os.path.join(fp_output, 'missing_emission_factors.csv'), index_col=[0])

    countries_missing_ef = dict((k, list(v.dropna())) for k, v in missing_factors.iteritems())  # make dict of tec: [countries]
    countries_missing_ef = {key:value for key, value in countries_missing_ef.items() if len(value)}  # remove entries with empty values

    """ Fix missing emission factors """
    # Note that these disregard whether or not there is actual production
    ef['Fossil Oil shale'] = ef['Fossil Oil'] # Proxy shale oil with conventional oil
    ef['Waste'] = 0  # Allocate incineration emissions to waste treatment as per ecoinvent
    ef['Marine'] = 0

    # Put in proxy for missing regions and technologies
    missing_factors.dropna(how='all', axis=1, inplace=True)
    missing_factors_dict = {}
    # build dict of missing technology-country factors
    for tec, col in missing_factors.iteritems():
        missing_factors_dict[tec] = list(col.dropna(how='any', axis=0).values)

    # countries with no geo-specific factors
    missing_countries = list(set(gen_df.index) - set(ef.index))
    for country in missing_countries:
        ef = ef.append(pd.Series(name=country))

    temp_gen_df = gen_df.loc[missing_countries]
    temp_gen_df = temp_gen_df.replace(0, np.nan).dropna(how='all', axis=1)
    temp_dict = {}
    for tec, col in temp_gen_df.items():
        temp = temp_gen_df[tec] > 0
        temp_ind = temp.index
        temp_dict[tec] = temp_ind[temp.values].to_list()

    """ Merge dictionaries and keep values of common keys in list"""
    def mergeDict(dict1, dict2):
        dict3 = {**dict1, **dict2}
        for key, value in dict3.items():
                if key in dict1 and key in dict2:
                    dict3[key] = (dict1[key] + dict2[key])
        return dict3

    # Merge dictionaries and add values of common keys in a list
    missing_factors_dict2 = mergeDict(missing_factors_dict, temp_dict)
    missing_factors_dict2

    # Use continental arithmetic average of same technology
    for key, countries in missing_factors_dict2.items():
        if not key=='Other' and not key=='Other renewable': # these are dealt with below
            ef[key].loc[countries] = ef[key].mean()

    # Make 'other' and 'other renewable' approximations
    renewable_dict = {'renewable': ['Biomass', 'Geothermal', 'Hydro Pumped Storage', 'Hydro Run-of-river and poundage',
                                    'Hydro Water Reservoir','Solar','Waste','Wind Onshore', 'Wind Offshore','Marine'],
                    'non-renewable': ['Fossil Gas','Fossil Hard coal','Fossil Oil','Fossil Brown coal/Lignite','Nuclear','Fossil Coal-derived gas','Fossil Oil shale','Fossil Peat']}

    other_tec = gen_df.loc[:,renewable_dict['non-renewable']]
    other_wf = other_tec.div(other_tec.sum(axis=1), axis=0) # determine generation shares for each country

    other_renew = gen_df.loc[:,renewable_dict['renewable']]
    other_renew_wf = other_renew.div(other_renew.sum(axis=1), axis=0)

    ef['Other'] = (ef*other_wf).sum(axis=1)
    ef['Other renewable'] = (ef*other_renew_wf).sum(axis=1)

    # Sort the indices of production, trade and emission factor matrices before export
    def sort_indices(df):
        df.sort_index(axis=0, inplace=True)
        df.sort_index(axis=1, inplace=True)

    sort_indices(trade_df)
    sort_indices(gen_df)
    sort_indices(ef)

    # Export all data
    # .csv for use in BEV_footprints_calculations.py

    ef.to_csv(os.path.join(fp_output,'final_emission_factors_' + str(year) + '.csv'))

    with pd.ExcelWriter(os.path.join(fp_results, 'entsoe_export_final_' + str(year) + '.xlsx')) as writer:
        trade_df.to_excel(writer, 'trade')
        gen_df.to_excel(writer, 'generation')
        ef.to_excel(writer,'new ef')

    return countries_missing_ef


import math
import sys
import os

from config import cfg
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.backends import cudnn
from datasets.GTA5.utils_GTA import colorize_mask
from models.fcn8_ssd import FCN8ResNet,FCN8VGG
from models.dc import dc_pixel,dc_object,dc_object2
from datasets.encoder import DataEncoder
import torch.optim as optim
from PIL import ImageDraw

import pdb

def init_model(key):

    resume = cfg.TRAIN.RESUME
    etractNet=[]
    dcNet=[]

    if key=='res':
        etractNet = FCN8ResNet(num_classes=cfg.DATA.NUM_CLASSES)
        dcNet = dc_pixel(input_channels=2048)
    elif key=='dense':
        etractNet = FCN8DenseNet(num_classes=cfg.DATA.NUM_CLASSES)
        dcNet = dc_pixel(input_channels=1920)
    elif key=='vgg':
        etractNet = FCN8VGG(num_classes=cfg.DATA.NUM_CLASSES)
        dcNet = dc_pixel(input_channels=512)
    else:
        print 'Invalid Networks!!!'
    obc_Net = dc_object(num_classes=cfg.DATA.NUM_CLASSES)

    max_epoch = 0
    if resume:
        ckpt_path = './ckpt/' + cfg.TRAIN.EXP_NAME + '/'
        extName,max_epoch = find_newest_model(ckpt_path)
        if os.path.exists(ckpt_path + extName):
            etractNet.load_state_dict(torch.load(ckpt_path + extName))
            print 'Successfully load ' + extName + '.'
        else:
            print 'No ext resume!!!'

        if os.path.exists(ckpt_path + 'dc.pth'):
            dcNet.load_state_dict(torch.load(ckpt_path + 'dc.pth'))
            print 'Successfully load dc.pth.'
        else:
            print 'No dc resume!!!'

        if os.path.exists(ckpt_path + 'obc.pth'):
            obc_Net.load_state_dict(torch.load(ckpt_path + 'obc.pth'))
            print 'Successfully load dc.pth.'
        else:
            print 'No obc resume!!!'

    if cfg.TRAIN.MULTI_GPU:
        etractNet = torch.nn.DataParallel(etractNet, device_ids=cfg.TRAIN.GPU_ID).cuda()
        dcNet = torch.nn.DataParallel(dcNet, device_ids=cfg.TRAIN.GPU_ID).cuda()
        obc_Net = torch.nn.DataParallel(obc_Net, device_ids=cfg.TRAIN.GPU_ID).cuda()
    else:
        etractNet = etractNet.cuda()
        dcNet = dcNet.cuda()
        obc_Net = obc_Net.cuda()

    cudnn.benchmark = True

    etractNet.train()
    dcNet.train()
    obc_Net.train()

    return etractNet, dcNet, obc_Net, max_epoch

def find_newest_model(path):
    epoch = []
    for model_name in os.listdir(path):
        if model_name.endswith('_ext.pth'):
            i_epoch = int(model_name.split('_')[1])
            epoch.append(i_epoch)
    max_epoch = max(epoch)
    newest_model = 'epoch_'+ str(max_epoch) + '_ext.pth'

    return newest_model,max_epoch

def set_extract_optimizer(etractNet, i_epoch):
    pretrained_lr = cfg.TRAIN.EXT_LR
    new_lr = cfg.TRAIN.SEG_LR
    ssd_lr = cfg.TRAIN.SSD_LR
    weight_decay = cfg.TRAIN.WEIGHT_DECAY

    ssd_lr = ssd_lr * (cfg.TRAIN.SSD_LR_DECAY_RATE **(i_epoch / cfg.TRAIN.SSD_LR_DECAY_EPOCH))
    new_lr = new_lr * (cfg.TRAIN.SEG_LR_DECAY_RATE **(i_epoch / cfg.TRAIN.SEG_LR_DECAY_EPOCH))
    extracter_optimizer = optim.SGD([
        {'params': [param for name, param in etractNet.named_parameters() if
                    name[-4:] == 'bias' and 'fconv' in name], 'lr': 2 * new_lr},
        {'params': [param for name, param in etractNet.named_parameters() if
                    name[-4:] != 'bias' and 'fconv' in name], 'lr': new_lr, 'weight_decay': weight_decay},
        {'params': [param for name, param in etractNet.named_parameters() if
                    name[-4:] == 'bias' and 'ssd' in name], 'lr': 2 * ssd_lr},
        {'params': [param for name, param in etractNet.named_parameters() if
                    name[-4:] != 'bias' and 'ssd' in name],'lr': ssd_lr, 'weight_decay': weight_decay},
        {'params': [param for name, param in etractNet.named_parameters() if
                    name[-4:] == 'bias' and 'fconv' not in name and 'ssd' not in name], 'lr': 2 * pretrained_lr},
        {'params': [param for name, param in etractNet.named_parameters() if
                    name[-4:] != 'bias' and 'fconv' not in name and 'ssd' not in name],'lr': pretrained_lr, 'weight_decay': weight_decay},

    ], momentum=0.9, nesterov=True)
    return extracter_optimizer

def calculate_mean_iu(predictions, gts):
    sum_iu = 0
    class_iu = np.zeros([cfg.DATA.NUM_CLASSES])
    for i in xrange(cfg.DATA.NUM_CLASSES):
        n_ii = t_i = sum_n_ji = 1e-9
        for p, gt in zip(predictions, gts):
            n_ii += np.sum(gt[p == i] == i)
            t_i += np.sum(gt == i)
            sum_n_ji += np.sum(p == i)

        class_iu[i] = float(n_ii) / (t_i + sum_n_ji - n_ii)
        sum_iu += float(n_ii) / (t_i + sum_n_ji - n_ii)
    mean_iu = sum_iu / (cfg.DATA.NUM_CLASSES)
    return mean_iu, class_iu

def logger(i_epoch,i_iter,i_tb,writer,ext_loss,sp_loss,unsp_loss,det_loss,src_m_iu,tgt_m_iu,
            dc_loss=None,dc_ins_loss=None,
            obc_loss=None,obc_ins_loss=None):

    batch_size = cfg.TRAIN.IMG_BATCH_SIZE
    writer.add_scalar('loss_ext', ext_loss.data[0], i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)
    writer.add_scalar('loss_sp', sp_loss.data[0], i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)
    writer.add_scalar('loss_unsp', unsp_loss.data[0], i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)
    writer.add_scalar('loss_det', det_loss.data[0], i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)
    writer.add_scalar('meanIU_src', src_m_iu, i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)
    writer.add_scalar('meanIU_tgt', tgt_m_iu, i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)

    if dc_loss is not None:
        writer.add_scalar('loss_dc', dc_loss.data[0], i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)
        writer.add_scalar('loss_dc_ins', dc_ins_loss.data[0], i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)

    if obc_loss is not None:
        writer.add_scalar('loss_obc', obc_loss.data[0], i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)
        writer.add_scalar('loss_obc_ins', obc_ins_loss.data[0], i_tb*batch_size*cfg.TRAIN.PRINT_FREQ)

    # print infos in terminial
    if cfg.TRAIN.COM_EXP == 5:
        print '[ep %d], [it %d], [Loss: dc %.4f, ext_all %.4f, sp %.4f, det %.4f, ivs1 %.4f, ivs2 %.4f], [src_miu %.4f], [tgt_miu %.4f]' \
            % (i_epoch + 1, i_iter + 1, dc_loss.data[0], ext_loss.data[0], sp_loss.data[0], det_loss.data[0], \
            dc_ins_loss.data[0], obc_ins_loss.data[0],src_m_iu,tgt_m_iu)

    elif cfg.TRAIN.COM_EXP == 4:
        print '[ep %d], [it %d], [Loss: dc %.4f, ext_all %.4f, sp %.4f, det %.4f, ivs1 %.4f], [src_miu %.4f], [tgt_miu %.4f]' \
            % (i_epoch + 1, i_iter + 1, dc_loss.data[0], ext_loss.data[0], sp_loss.data[0], det_loss.data[0], \
                dc_ins_loss.data[0], src_m_iu,tgt_m_iu)

    elif cfg.TRAIN.COM_EXP == 3:
        print '[ep %d], [it %d], [Loss: ext_all %.4f, sp %.4f, det %.4f], [src_miu %.4f], [tgt_miu %.4f]' \
        % (i_epoch + 1, i_iter + 1, ext_loss.data[0], sp_loss.data[0], det_loss.data[0], src_m_iu,tgt_m_iu)

    elif cfg.TRAIN.COM_EXP == 6:
        print '[ep %d], [it %d], [Loss: ext_all %.4f, sp %.4f, det %.4f, ivs2 %.4f], [src_miu %.4f], [tgt_miu %.4f]' \
            % (i_epoch + 1, i_iter + 1, ext_loss.data[0], sp_loss.data[0], det_loss.data[0], \
         obc_ins_loss.data[0],src_m_iu,tgt_m_iu)

def show_img(i_tb, writer,src_label_color,src_pred_color,tgt_label_color,tgt_pred_color,
            src_det_img, tgt_det_img):

    batch_size = cfg.TRAIN.IMG_BATCH_SIZE
    writer.add_image('gt_src', src_label_color)
    writer.add_image('pred_src', src_pred_color)
    writer.add_image('det_src', np.array(src_det_img))

    writer.add_image('gt_tgt', tgt_label_color)
    writer.add_image('pred_tgt', tgt_pred_color)
    writer.add_image('det_tgt', np.array(tgt_det_img))

def eval_ext_logs(outputs,labels):
    outputs = outputs[:, :cfg.DATA.NUM_CLASSES - 1, :, :]
    # pdb.set_trace()

    prediction = outputs.data.max(1)[1].squeeze_(1).cpu().numpy()
    mean_iu, class_iu = calculate_mean_iu(prediction, labels.data.cpu().numpy())

    return mean_iu, class_iu

def gen_labeled_map(outputs,labels=None):
    map_size = cfg.TRAIN.IMG_SIZE
    dst_label_shape = (map_size[0],map_size[1],1)
    dst_RGB_shape = (map_size[0],map_size[1],3)
    # pdb.set_trace()
    prediction = outputs.data.max(1)[1].squeeze_(1).cpu().numpy()
    pred_colormap = np.array(colorize_mask(prediction[0])).reshape(dst_RGB_shape)
    if labels is None:
    	return pred_colormap
    # pdb.set_trace()
    label_colormap = np.array(colorize_mask(labels.data[0].cpu().numpy())).reshape(dst_RGB_shape)

    return pred_colormap, label_colormap

# Done
def draw_bbx(loc_preds, conf_preds, img, restore):

    data_encoder = DataEncoder()
    bbx, _, _ = \
    data_encoder.decode(loc_preds.data.squeeze(0), F.softmax(conf_preds.squeeze(0)).data)

    map_size = cfg.TRAIN.IMG_SIZE
    bbx = bbx * map_size[0]
    # pdb.set_trace()
    img = restore(img)
    # print img.size()
    imgDraw = ImageDraw.Draw(img)
    for i_bbx in range(0,bbx.shape[0]):
        imgDraw.rectangle(bbx[i_bbx,:].tolist(),outline = "red")
    del imgDraw

    return img

def gen_obc_labels(src_obj_labels,tgt_obj_labels):
    len_src_obj = src_obj_labels.size()[0]*src_obj_labels.size()[1]
    len_tgt_obj = tgt_obj_labels.size()[0]*tgt_obj_labels.size()[1]

    src_obj_labels = src_obj_labels.resize_(len_src_obj,1)
    tgt_obj_labels = tgt_obj_labels.resize_(len_tgt_obj,1)

    obc_labels = torch.zeros((src_obj_labels.size()[0] + tgt_obj_labels.size()[0]))
    obc_ivs_labels = torch.zeros(obc_labels.shape)

    for i_src, label in enumerate(src_obj_labels):
        obc_labels[i_src] = label.numpy()[0]
        obc_ivs_labels[i_src] = cfg.DATA.NUM_CLASSES + label.numpy()[0]
    for i_tgt, label in enumerate(tgt_obj_labels):
        obc_labels[src_obj_labels.size()[0] + i_tgt] = cfg.DATA.NUM_CLASSES + label.numpy()[0]
        obc_ivs_labels[src_obj_labels.size()[0] + i_tgt] = label.numpy()[0]

    obc_labels = Variable(obc_labels.long().cuda())
    obc_ivs_labels = Variable(obc_ivs_labels.long().cuda())
    return obc_labels, obc_ivs_labels

def gen_dc_pixels_label():
    batch_size = cfg.TRAIN.IMG_BATCH_SIZE
    dc_labels = np.zeros((2 * batch_size, cfg.TRAIN.IMG_SIZE[0],cfg.TRAIN.IMG_SIZE[1]))
    dc_ivs_labels = np.zeros(dc_labels.shape)
    dc_labels[:batch_size,:,:] = 1
    dc_ivs_labels[batch_size:,:,:] = 1
    dc_labels = torch.from_numpy(dc_labels.astype(np.int64))
    dc_labels = Variable(dc_labels.cuda())
    dc_ivs_labels = torch.from_numpy(dc_ivs_labels.astype(np.int64))
    dc_ivs_labels = Variable(dc_ivs_labels.cuda())
    return dc_labels, dc_ivs_labels

def forward_ext_model(ext_model, ext_src_inputs, ext_tgt_inputs,
                        src_bbx,tgt_bbx):
    pred_tgt_outputs, dc_tgt_inputs, tgt_loc_preds, tgt_conf_preds, tgt_pooled_features \
                        = ext_model(Variable(ext_tgt_inputs.cuda()),gt=Variable(tgt_bbx.cuda()))

    pred_src_outputs, dc_src_inputs, src_loc_preds, src_conf_preds, src_pooled_features \
                    = ext_model(Variable(ext_src_inputs.cuda()),gt=Variable(src_bbx.cuda()))

    loc_preds = torch.cat((src_loc_preds, tgt_loc_preds),0)
    conf_preds = torch.cat((src_conf_preds, tgt_conf_preds),0)
    dc_inputs = torch.cat((dc_src_inputs, dc_tgt_inputs),0)
    pooled_features = torch.cat((src_pooled_features, tgt_pooled_features),0)
    return pred_src_outputs, pred_tgt_outputs, loc_preds, conf_preds, dc_inputs, pooled_features

def preapre_bbx_for_roipol(bbx):
    len_obj = bbx.size()[0]*bbx.size()[1]

    img_batch_size = cfg.TRAIN.IMG_BATCH_SIZE
    obj_batch_size = cfg.TRAIN.OB_BATCH_SIZE

    if len_obj != obj_batch_size*img_batch_size:
        print 'Batch Size is ERROR!'

    bbx = bbx.resize_(len_obj,4)

    batch_ind = torch.zeros(len_obj,1)
    for i_img in range(0,img_batch_size):
        batch_ind[i_img*obj_batch_size:(i_img+1)*obj_batch_size,0] = i_img
    # pdb.set_trace()
    bbx = torch.cat((batch_ind,bbx),1)
    return bbx


# Sweep animation
# Moves a dot from one side to the other, one pixel at a time.

import sys, traceback, random
from numpy import array,full

class strip_animation():
    def __init__(self,datastore):
        self.max_led=datastore.LED_COUNT
        self.lamp_length=datastore.LAMP_LENGTH
        self.brightness_scaling=int(255/self.lamp_length)
        self.sweep_pos=0
        self.colpos=0
        self.nextpos=0
        self.changestrip=0

    def emit_row(self):
        try:
            if self.changestrip == 0 and random.randint(0,1000) == 800:
                self.nextpos=random.randint(0,3)
                if self.nextpos != self.colpos:
                    self.changestrip=1
                    #print("Cur Col: ", self.colpos, "Next: ",self.nextpos)
            if random.randint(0,10)==5:
                self.sweep_pos = (self.sweep_pos +1) % (self.max_led-1)
            row_arr=full((self.max_led,4),0)
            brt=((self.sweep_pos % self.lamp_length)*self.brightness_scaling)
            curbg=240-((240*self.changestrip)>>8)
            if self.changestrip > 0:
                newbg=(240*self.changestrip)>>8
                row_arr[(((int(self.sweep_pos/self.lamp_length)-3)*self.lamp_length)) % self.max_led][self.nextpos]=newbg
                row_arr[(((int(self.sweep_pos/self.lamp_length)-2)*self.lamp_length)) % self.max_led][self.nextpos]=newbg
                row_arr[(((int(self.sweep_pos/self.lamp_length)-1)*self.lamp_length)) % self.max_led][self.nextpos]=(brt*self.changestrip)>>8
                row_arr[int(self.sweep_pos/self.lamp_length)*self.lamp_length][self.colpos]=0
                row_arr[(((int(self.sweep_pos/self.lamp_length)+1)*self.lamp_length)) % self.max_led][self.nextpos]=newbg-((brt*self.changestrip)>>8)
                row_arr[(((int(self.sweep_pos/self.lamp_length)+2)*self.lamp_length)) % self.max_led][self.nextpos]=newbg
                row_arr[(((int(self.sweep_pos/self.lamp_length)+3)*self.lamp_length)) % self.max_led][self.nextpos]=newbg
                #print("str",self.changestrip,"btr",brt,"nbg",newbg,"nlo",(brt*self.changestrip)>>8,"nhi",newbg-((brt*self.changestrip)>>8),"cbg",curbg,"clo",brt-((brt*self.changestrip)>>8),"chi",curbg-(brt-((brt*self.changestrip)>>8)))

                if self.changestrip < 256:
                    self.changestrip += 1
                else:
                    self.changestrip = 0
                    curbg=240-((240*self.changestrip)>>8)
                    self.colpos = self.nextpos
            row_arr[(((int(self.sweep_pos/self.lamp_length)-3)*self.lamp_length)) % self.max_led][self.colpos]=curbg
            row_arr[(((int(self.sweep_pos/self.lamp_length)-2)*self.lamp_length)) % self.max_led][self.colpos]=curbg
            row_arr[(((int(self.sweep_pos/self.lamp_length)-1)*self.lamp_length)) % self.max_led][self.colpos]=brt-((brt*self.changestrip)>>8)
            row_arr[int(self.sweep_pos/self.lamp_length)*self.lamp_length][self.colpos]=0
            row_arr[(((int(self.sweep_pos/self.lamp_length)+1)*self.lamp_length)) % self.max_led][self.colpos]=curbg-(brt-((brt*self.changestrip)>>8))
            row_arr[(((int(self.sweep_pos/self.lamp_length)+2)*self.lamp_length)) % self.max_led][self.colpos]=curbg
            row_arr[(((int(self.sweep_pos/self.lamp_length)+3)*self.lamp_length)) % self.max_led][self.colpos]=curbg
            return row_arr
        except Exception as err:
            print(err)
            traceback.print_exc(file=sys.stdout)


import torch
import pit
from deit.losses import DistillationLoss
import numpy as np
from reprod_log import ReprodLogger
from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
import argparse

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--smoothing', type=float, default=0.1, help='Label smoothing (default: 0.1)')
parser.add_argument('--mixup', type=float, default=0.8,
                    help='mixup alpha, mixup enabled if > 0. (default: 0.8)')
args = parser.parse_args()
model = pit.pit_ti(pretrained=False)
model.load_state_dict(torch.load('weights/pit_ti_730.pth',map_location='cpu'))
fake_data = np.load("fake_data.npy")
fake_label = np.load("fake_label_int.npy")
teacher_model = None
criterion = LabelSmoothingCrossEntropy()
if args.mixup > 0.:
    # smoothing is handled with mixup label transform
    criterion = SoftTargetCrossEntropy()
elif args.smoothing:
    criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
else:
    criterion = torch.nn.CrossEntropyLoss()
dis = DistillationLoss(
    criterion, teacher_model, None, 0.5, 1.0
)
out = model(torch.tensor(fake_data))
loss = dis(torch.tensor(fake_data),out,torch.tensor(fake_label))
reprod_logger = ReprodLogger()
reprod_logger.add("loss", loss.cpu().detach().numpy())
reprod_logger.save("loss_torch.npy")


# -*- encoding: utf-8 -*-

import numpy as np
import six
import warnings

import autosklearn.automl
from autosklearn.constants import *
from autosklearn.util.backend import create
from sklearn.base import BaseEstimator
import sklearn.utils
import scipy.sparse

class AutoMLDecorator(object):

    def __init__(self, automl):
        self._automl = automl

    def fit(self, *args, **kwargs):
        self._automl.fit(*args, **kwargs)

    def refit(self, X, y):
        """Refit all models found with fit to new data.

        Necessary when using cross-validation. During training, auto-sklearn
        fits each model k times on the dataset, but does not keep any trained
        model and can therefore not be used to predict for new data points.
        This methods fits all models found during a call to fit on the data
        given. This method may also be used together with holdout to avoid
        only using 66% of the training data to fit the final model.

        Parameters
        ----------

        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The training input samples.

        y : array-like, shape = [n_samples] or [n_samples, n_outputs]
            The targets.

        Returns
        -------

        self

        """
        return self._automl.refit(X, y)

    def fit_ensemble(self, y, task=None, metric=None, precision='32',
                     dataset_name=None, ensemble_nbest=None,
                     ensemble_size=None):
        """Build the ensemble.

        This method only needs to be called in the parallel mode.

        Returns
        -------
        self
        """
        return self._automl.fit_ensemble(y, task, metric, precision,
                                         dataset_name, ensemble_nbest,
                                         ensemble_size)

    def predict(self, X):
        return self._automl.predict(X)

    def score(self, X, y):
        return self._automl.score(X, y)

    def show_models(self):
        """Return a representation of the final ensemble found by auto-sklearn

        Returns
        -------
        str
        """
        return self._automl.show_models()

    @property
    def grid_scores_(self):
        return self._automl.grid_scores_

    @property
    def cv_results_(self):
        return self._automl.cv_results_

    @property
    def trajectory_(self):
        return self._automl.trajectory_

    def sprint_statistics(self):
        return self._automl.sprint_statistics()

class AutoSklearnEstimator(AutoMLDecorator, BaseEstimator):

    def __init__(self,
                 time_left_for_this_task=3600,
                 per_run_time_limit=360,
                 initial_configurations_via_metalearning=25,
                 ensemble_size=50,
                 ensemble_nbest=50,
                 seed=1,
                 ml_memory_limit=3072,
                 include_estimators=None,
                 exclude_estimators=None,
                 include_preprocessors=None,
                 exclude_preprocessors=None,
                 resampling_strategy='holdout',
                 resampling_strategy_arguments=None,
                 tmp_folder=None,
                 output_folder=None,
                 delete_tmp_folder_after_terminate=True,
                 delete_output_folder_after_terminate=True,
                 shared_mode=False,
                 disable_evaluator_output=False,
                 configuration_mode='SMAC'):
        """
        Parameters
        ----------
        time_left_for_this_task : int, optional (default=3600)
            Time limit in seconds for the search of appropriate
            models. By increasing this value, *auto-sklearn* has a higher
            chance of finding better models.

        per_run_time_limit : int, optional (default=360)
            Time limit for a single call to the machine learning model.
            Model fitting will be terminated if the machine learning
            algorithm runs over the time limit. Set this value high enough so
            that typical machine learning algorithms can be fit on the
            training data.

        initial_configurations_via_metalearning : int, optional (default=25)
            Initialize the hyperparameter optimization algorithm with this
            many configurations which worked well on previously seen
            datasets. Disable if the hyperparameter optimization algorithm
            should start from scratch.

        ensemble_size : int, optional (default=50)
            Number of models added to the ensemble built by `Ensemble
            selection from libraries of models. Models are drawn with
            replacement.

        ensemble_nbest : int, optional (default=50)
            Only consider the ``ensemble_nbest`` models when building an
            ensemble. Implements `Model Library Pruning` from `Getting the
            most out of ensemble selection`.

        seed : int, optional (default=1)

        ml_memory_limit : int, optional (3072)
            Memory limit in MB for the machine learning algorithm.
            `auto-sklearn` will stop fitting the machine learning algorithm if
            it tries to allocate more than `ml_memory_limit` MB.

        include_estimators : list, optional (None)
            If None, all possible estimators are used. Otherwise specifies
            set of estimators to use.

        exclude_estimators : list, optional (None)
            If None, all possible estimators are used. Otherwise specifies
            set of estimators not to use. Incompatible with include_estimators.

        include_preprocessors : list, optional (None)
            If None all possible preprocessors are used. Otherwise specifies set
            of preprocessors to use.

        exclude_preprocessors : list, optional (None)
            If None all possible preprocessors are used. Otherwise specifies set
            of preprocessors not to use. Incompatible with
            include_preprocessors.

        resampling_strategy : string, optional ('holdout')
            how to to handle overfitting, might need 'resampling_strategy_arguments'

            * 'holdout': 66:33 (train:test) split
            * 'holdout-iterative-fit':  66:33 (train:test) split, calls iterative
              fit where possible
            * 'cv': crossvalidation, requires 'folds'

        resampling_strategy_arguments : dict, optional if 'holdout' (None)
            Additional arguments for resampling_strategy
            * 'holdout': None
            * 'holdout-iterative-fit':  None
            * 'cv': {'folds': int}

        tmp_folder : string, optional (None)
            folder to store configuration output and log files, if ``None``
            automatically use ``/tmp/autosklearn_tmp_$pid_$random_number``

        output_folder : string, optional (None)
            folder to store predictions for optional test set, if ``None``
            automatically use ``/tmp/autosklearn_output_$pid_$random_number``

        delete_tmp_folder_after_terminate: string, optional (True)
            remove tmp_folder, when finished. If tmp_folder is None
            tmp_dir will always be deleted

        delete_output_folder_after_terminate: bool, optional (True)
            remove output_folder, when finished. If output_folder is None
            output_dir will always be deleted

        shared_mode: bool, optional (False)
            Run smac in shared-model-node. This only works if arguments
            ``tmp_folder`` and ``output_folder`` are given and both
            ``delete_tmp_folder_after_terminate`` and
            ``delete_output_folder_after_terminate`` are set to False.

        disable_evaluator_output: bool, optional (False)
            Disable model and prediction output. Cannot be used together with
            ensemble building. predict() cannot be used when setting this
            flag to True.

        Attributes
        ----------
        grid_scores\_ : list of named tuples
            Contains scores for all parameter combinations in param_grid.
            Each entry corresponds to one parameter setting.
            Each named tuple has the attributes:

            * ``parameters``, a dict of parameter settings
            * ``mean_validation_score``, the mean score over the
              cross-validation folds
            * ``cv_validation_scores``, the list of scores for each fold

        cv_results\_ : dict of numpy (masked) ndarrays
            A dict with keys as column headers and values as columns, that can be
            imported into a pandas ``DataFrame``.

            This attribute is a backward port to already support the advanced
            output of scikit-learn 0.18. Not all keys returned by scikit-learn
            are supported yet.

        """
        self.time_left_for_this_task = time_left_for_this_task
        self.per_run_time_limit = per_run_time_limit
        self.initial_configurations_via_metalearning = initial_configurations_via_metalearning
        self.ensemble_size = ensemble_size
        self.ensemble_nbest = ensemble_nbest
        self.seed = seed
        self.ml_memory_limit = ml_memory_limit
        self.include_estimators = include_estimators
        self.exclude_estimators = exclude_estimators
        self.include_preprocessors = include_preprocessors
        self.exclude_preprocessors = exclude_preprocessors
        self.resampling_strategy = resampling_strategy
        self.resampling_strategy_arguments = resampling_strategy_arguments
        self.tmp_folder = tmp_folder
        self.output_folder = output_folder
        self.delete_tmp_folder_after_terminate = delete_tmp_folder_after_terminate
        self.delete_output_folder_after_terminate = delete_output_folder_after_terminate
        self.shared_mode = shared_mode
        self.disable_evaluator_output = disable_evaluator_output
        self.configuration_mode = configuration_mode
        super(AutoSklearnEstimator, self).__init__(None)

    def build_automl(self):
        if self.shared_mode:
            self.delete_output_folder_after_terminate = False
            self.delete_tmp_folder_after_terminate = False
            if self.tmp_folder is None:
                raise ValueError("If shared_mode == True tmp_folder must not "
                                 "be None.")
            if self.output_folder is None:
                raise ValueError("If shared_mode == True output_folder must "
                                 "not be None.")

        backend = create(temporary_directory=self.tmp_folder,
                         output_directory=self.output_folder,
                         delete_tmp_folder_after_terminate=self.delete_tmp_folder_after_terminate,
                         delete_output_folder_after_terminate=self.delete_output_folder_after_terminate)
        automl = autosklearn.automl.AutoML(
            backend=backend,
            time_left_for_this_task=self.time_left_for_this_task,
            per_run_time_limit=self.per_run_time_limit,
            log_dir=backend.temporary_directory,
            initial_configurations_via_metalearning=
            self.initial_configurations_via_metalearning,
            ensemble_size=self.ensemble_size,
            ensemble_nbest=self.ensemble_nbest,
            seed=self.seed,
            ml_memory_limit=self.ml_memory_limit,
            include_estimators=self.include_estimators,
            exclude_estimators=self.exclude_estimators,
            include_preprocessors=self.include_preprocessors,
            exclude_preprocessors=self.exclude_preprocessors,
            resampling_strategy=self.resampling_strategy,
            resampling_strategy_arguments=self.resampling_strategy_arguments,
            delete_tmp_folder_after_terminate=self.delete_tmp_folder_after_terminate,
            delete_output_folder_after_terminate=
            self.delete_output_folder_after_terminate,
            shared_mode=self.shared_mode,
            configuration_mode=self.configuration_mode,
            disable_evaluator_output=self.disable_evaluator_output)

        return automl

    def fit(self, *args, **kwargs):
        self._automl = self.build_automl()
        super(AutoSklearnEstimator, self).fit(*args, **kwargs)

    def fit_ensemble(self, y, task=None, metric=None, precision='32',
                     dataset_name=None, ensemble_nbest=None,
                     ensemble_size=None):
        if self._automl is None:
            self._automl = self.build_automl()
        return self._automl.fit_ensemble(y, task, metric, precision,
                                         dataset_name, ensemble_nbest,
                                         ensemble_size)

class AutoSklearnClassifier(AutoSklearnEstimator):
    """
    This class implements the classification task.
    """

    def build_automl(self):
        automl = super(AutoSklearnClassifier, self).build_automl()
        return AutoMLClassifier(automl)

    def fit(self, X, y,
            metric='acc_metric',
            feat_type=None,
            dataset_name=None):
        """Fit *auto-sklearn* to given training set (X, y).

        Parameters
        ----------

        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The training input samples.

        y : array-like, shape = [n_samples] or [n_samples, n_outputs]
            The target classes.

        metric : str, optional (default='acc_metric')
            The metric to optimize for. Can be one of: ['acc_metric',
            'auc_metric', 'bac_metric', 'f1_metric', 'pac_metric']. A
            description of the metrics can be found in `the paper describing
            the AutoML Challenge.
        feat_type : list, optional (default=None)
            List of str of `len(X.shape[1])` describing the attribute type.
            Possible types are `Categorical` and `Numerical`. `Categorical`
            attributes will be automatically One-Hot encoded.

        dataset_name : str, optional (default=None)
            Create nicer output. If None, a string will be determined by the
            md5 hash of the dataset.

        Returns
        -------
        self

        """
        return super(AutoSklearnClassifier, self).fit(X=X, y=y, metric=metric,
                                                      feat_type=feat_type,
                                                      dataset_name=dataset_name)

    def predict(self, X):
        """Predict classes for X.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]

        Returns
        -------
        y : array of shape = [n_samples] or [n_samples, n_labels]
            The predicted classes.

        """
        return super(AutoSklearnClassifier, self).predict(X)

    def predict_proba(self, X):

        """Predict probabilities of classes for all samples X.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]

        Returns
        -------
        y : array of shape = [n_samples, n_classes] or [n_samples, n_labels]
            The predicted class probabilities.
        """
        return self._automl.predict_proba(X)

class AutoSklearnRegressor(AutoSklearnEstimator):
    """
    This class implements the regression task.
    """

    def build_automl(self):
        automl = super(AutoSklearnRegressor, self).build_automl()
        return AutoMLRegressor(automl)

    def fit(self, X, y,
            metric='r2_metric',
            feat_type=None,
            dataset_name=None):
        """Fit *autosklearn* to given training set (X, y).

        Parameters
        ----------

        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The training input samples.

        y : array-like, shape = [n_samples] or [n_samples, n_outputs]
            The regression target.

        metric : str, optional (default='r2_metric')
            The metric to optimize for. Can be one of: ['r2_metric',
            'a_metric']. A description of the metrics can be found in
            `the paper describing the AutoML Challenge
            <http://www.causality.inf.ethz.ch/AutoML/automl_ijcnn15.pdf>`_.

        feat_type : list, optional (default=None)
            List of str of `len(X.shape[1])` describing the attribute type.
            Possible types are `Categorical` and `Numerical`. `Categorical`
            attributes will be automatically One-Hot encoded.

        dataset_name : str, optional (default=None)
            Create nicer output. If None, a string will be determined by the
            md5 hash of the dataset.

        Returns
        -------
        self

        """
        # Fit is supposed to be idempotent!
        # But not if we use share_mode.
        return super(AutoSklearnRegressor, self).fit(X=X, y=y, metric=metric,
                                                     feat_type=feat_type,
                                                     dataset_name=dataset_name)

    def predict(self, X):
        """Predict regression target for X.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]

        Returns
        -------
        y : array of shape = [n_samples] or [n_samples, n_outputs]
            The predicted values.

        """
        return super(AutoSklearnRegressor, self).predict(X)

class AutoMLClassifier(AutoMLDecorator):

    def __init__(self, automl):
        self._classes = []
        self._n_classes = []
        self._n_outputs = 0

        super(AutoMLClassifier, self).__init__(automl)

    def fit(self, X, y,
            metric='acc_metric',
            loss=None,
            feat_type=None,
            dataset_name=None,
            ):
        # From sklearn.tree.DecisionTreeClassifier
        X = sklearn.utils.check_array(X, accept_sparse="csr",
                                      force_all_finite=False)
        if scipy.sparse.issparse(X):
            X.sort_indices()

        y = self._process_target_classes(y)

        if self._n_outputs > 1:
            task = MULTILABEL_CLASSIFICATION
        else:
            if len(self._classes[0]) == 2:
                task = BINARY_CLASSIFICATION
            else:
                task = MULTICLASS_CLASSIFICATION

        return self._automl.fit(X, y, task, metric, feat_type, dataset_name)

    def fit_ensemble(self, y, task=None, metric=None, precision='32',
                     dataset_name=None, ensemble_nbest=None,
                     ensemble_size=None):
        self._process_target_classes(y)
        return self._automl.fit_ensemble(y, task, metric, precision, dataset_name,
                                         ensemble_nbest, ensemble_size)

    def _process_target_classes(self, y):
        y = np.atleast_1d(y)
        if y.ndim == 2 and y.shape[1] == 1:
            warnings.warn("A column-vector y was passed when a 1d array was"
                          " expected. Please change the shape of y to "
                          "(n_samples,), for example using ravel().",
                          sklearn.utils.DataConversionWarning, stacklevel=2)

        if y.ndim == 1:
            # reshape is necessary to preserve the data contiguity against vs
            # [:, np.newaxis] that does not.
            y = np.reshape(y, (-1, 1))

        self._n_outputs = y.shape[1]

        y = np.copy(y)

        self._classes = []
        self._n_classes = []

        for k in six.moves.range(self._n_outputs):
            classes_k, y[:, k] = np.unique(y[:, k], return_inverse=True)
            self._classes.append(classes_k)
            self._n_classes.append(classes_k.shape[0])

        self._n_classes = np.array(self._n_classes, dtype=np.int)

        # TODO: fix metafeatures calculation to allow this!
        if y.shape[1] == 1:
            y = y.flatten()

        return y

    def predict(self, X):
        predicted_probabilities = self._automl.predict(X)
        if self._n_outputs == 1:
            predicted_indexes = np.argmax(predicted_probabilities, axis=1)
            predicted_classes = self._classes[0].take(predicted_indexes)

            return predicted_classes
        else:
            argmax_v = np.vectorize(np.argmax, otypes=[int])
            predicted_indexes = argmax_v(predicted_probabilities)
            #predicted_indexes = np.argmax(predicted_probabilities, axis=1)
            n_samples = predicted_probabilities.shape[0]
            predicted_classes = np.zeros((n_samples, self._n_outputs), dtype=object)

            for k in six.moves.range(self._n_outputs):
                output_predicted_indexes = predicted_indexes[:, k].reshape(-1)
                predicted_classes[:, k] = self._classes[k].take(output_predicted_indexes)

            return predicted_classes

    def predict_proba(self, X):
        return self._automl.predict(X)

class AutoMLRegressor(AutoMLDecorator):

    def fit(self, X, y,
            metric='r2_metric',
            loss=None,
            feat_type=None,
            dataset_name=None,
            ):
        return self._automl.fit(X=X, y=y, task=REGRESSION, metric=metric,
                                feat_type=feat_type, dataset_name=dataset_name)


from PIL import Image, ImageOps
import numpy as np
import skimage.io as io
from src.models.class_patcher import patcher
from src.utils.imgproc import *
from skimage.color import rgb2hsv, hsv2rgb

class patcher(patcher):
    def __init__(self, body='./body/body_glaze.png', **options):
        super().__init__('ぐれーず', body=body, pantie_position=[265, 2019], **options)
        self.mask = io.imread('./mask/mask_glaze.png')
        self.pantie_edge = io.imread('./material/glaze_pantie_edge.png') / 255.

        try:
            self.with_bra = self.options['with_bra']
        except:
            self.with_bra = self.ask(question='With bra?', default=True)
        if self.with_bra:
            self.bra_mask = io.imread('./mask/bra_glaze_center.png')
            self.bra_others = io.imread('./material/glaze_others.png') / 255.

    def convert(self, image):
        pantie = np.array(image)
        patch = np.copy(pantie[-140:-5, 546:, :])
        [pr, pc, d] = patch.shape
        arrx = np.zeros(25) - 10
        arry = np.zeros(25) - 5
        arrx[6:-10] += np.sin(np.linspace(0, np.pi, 9)) * 4
        arrx[-15:] -= np.sin(np.linspace(0, np.pi / 2, 15)) * 5
        arry[1:5] -= np.sin(np.linspace(0, np.pi, 4)) * 25
        arry[5:] += np.sin(np.linspace(0, np.pi / 2, 20)) * 40
        pantie = affine_transform_by_arr(pantie, arrx, arry)
        pantie[125:125 + pr, :pc, :] = np.float32(patch[::-1, ::-1] / 255)
        pantie = np.uint8(resize(pantie[:320, :-50], [3.25, 3.15]) * 255)[:, 38:]
        pantie = np.bitwise_and(pantie, self.mask)
        pantie = np.concatenate((pantie[:, ::-1], pantie), axis=1)
        return Image.fromarray(pantie)

    def gen_bra_center(self, image):
        pantie = np.array(image)
        center = np.uint8(resize(pantie[15:150, 450:], [7.85, 7.85]) * 255)
        center = np.bitwise_and(center, self.bra_mask)
        return Image.fromarray(center)

    def gen_colored_tex(self, image, texture):
        pantie = np.array(image)
        tex = np.copy(texture)
        front = pantie[20:100, 30:80, :3] / 255.0
        front_color = np.mean(np.mean(front, axis=0), axis=0)
        front_shade = pantie[130:150, 0:40, :3] / 255.0
        front_shade_color = np.mean(np.mean(front_shade, axis=0), axis=0)
        # tex[:, :, :3] = tex[:, :, :3] * front_shade_color
        tex[:, :, :3] = tex[:, :, :3] * front_color
        return Image.fromarray(np.uint8(tex * 255))

    def patch(self, image, transparent=False):
        pantie = self.convert(image)
        if transparent:
            patched = Image.new("RGBA", (4096, 4096))
        else:
            patched = self.body.copy()
        patched = self.paste(patched, pantie, self.pantie_position)
        patched = self.paste(patched, self.gen_colored_tex(image, self.pantie_edge), [258, 1907])
        if self.with_bra:
            bra_center = self.gen_bra_center(image)
            patched = self.paste(patched, bra_center, [528, 162])
            patched = self.paste(patched, ImageOps.mirror(bra_center), [2173, 162])
            patched = self.paste(patched, self.gen_colored_tex(image, self.bra_others), [251, 921])
        return patched


#!/usr/bin/env python
# -*- coding: utf-8 -*-

#  This software and supporting documentation are distributed by
#      Institut Federatif de Recherche 49
#      CEA/NeuroSpin, Batiment 145,
#      91191 Gif-sur-Yvette cedex
#      France
#
# This software is governed by the CeCILL-B license under
# French law and abiding by the rules of distribution of free software.
# You can  use, modify and/or redistribute the software under the
# terms of the CeCILL-B license as circulated by CEA, CNRS
# and INRIA at the following URL "http://www.cecill.info".
#
# As a counterpart to the access to the source code and  rights to copy,
# modify and redistribute granted by the license, users are provided only
# with a limited warranty  and the software's author,  the holder of the
# economic rights,  and the successive licensors  have only  limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using,  modifying and/or developing or reproducing the
# software by the user in light of its specific status of free software,
# that may mean  that it is complicated to manipulate,  and  that  also
# therefore means  that it is reserved for developers  and  experienced
# professionals having in-depth computer knowledge. Users are therefore
# encouraged to load and test the software's suitability as regards their
# requirements in conditions enabling the security of their systems and/or
# data to be ensured and,  more generally, to use and operate it in the
# same conditions as regards security.
#
# The fact that you are presently reading this means that you have had
# knowledge of the CeCILL-B license and that you accept its terms.

from __future__ import print_function

from __future__ import absolute_import
from soma import aims, aimsalgo
from optparse import OptionParser
import numpy, math
from six.moves import range

parser = OptionParser( description = 'Creates a cone or other mesh '
  'representing the TMS active zone at the specified target position along '
  'the head mesh' )
parser.add_option( '-i', '--input', dest='headmeshfile',
  help='input head mesh' )
parser.add_option( '-o', '--output', dest='conefile',
  help='output mesh' )
parser.add_option( '-s', '--stimpoint', dest='stimpoint', type='float',
  nargs=3, help='stimulation position in mm on head or cortex ' \
    '(3 float numbers)' )
parser.add_option( '-r', '--radius', dest='baseradius', type='float',
  help='stimulation cone base radius' )
parser.add_option( '-d', '--depth', dest='depth', type='float',
  help='stimulation cone depth' )
parser.add_option( '-n', '--normalradius', dest='normalradius', type='float',
  help='radius around stimulation point to average normal [optional, ' \
  'default=20] ' )
parser.add_option( '-m', '--normalmesh', dest='normalmesh',
  help='output mesh for the normal line (just for visualization purpose) ' \
    '[optional]' )
parser.add_option( '-t', '--targetlinemesh', dest='targetlinemesh',
  help='output mesh for the target-head projection line, which ' \
  'might be slightly different from the normal. If the target is specified ' \
  'on the head, this line is by defninition the normal. [optional] ' )
parser.add_option( '-b', '--brain', dest='brain',
  help='brain (or grey/white cortex) mask [optional, -z option needed]' )
parser.add_option( '-c', '--cortex', dest='cortexvalue', type='int',
  help='cortex value in grey/white mask image, [default: all non-zero]' )
parser.add_option( '-z', '--zone', dest='stimzone',
  help='stimulation zone output filename (.bck) [optional, -b option needed]' )
parser.add_option( '-w', '--writeresults', dest='writeresults',
  help='output text file where to write results (points and directions) ' \
  '[optional]' )

parser.set_defaults( normalradius = 20, cortexvalue=None )

(options, args) = parser.parse_args()

#print(options)
if options.headmeshfile is None or options.conefile is None \
  or options.stimpoint is None or options.baseradius is None \
  or options.depth is None \
  or (options.brain is None and options.stimzone is not None):
  parser.parse_args( [ '-h' ] )

targetpoint = aims.Point3df( options.stimpoint )
meshfile = options.headmeshfile
depth = options.depth
baseradius = options.baseradius
outfile = options.conefile
normalradius = options.normalradius
if normalradius < 1:
  normalradius = 1

headmesh = aims.read( meshfile )
# project stimpoint to the head mesh
stimpoint = aims.SurfaceManip.nearestPointToMesh( targetpoint, headmesh,
  10 )[0]
# find closest vertices to stimpoint
dist = numpy.array( [ (p - stimpoint).norm() for p in headmesh.vertex() ] )
normals = headmesh.normal()
snormals = [ normals[n] for n in numpy.where( dist <= normalradius )[0] ]
ntarget = aims.Point3df( numpy.mean( snormals, 0 ) )
ntarget.normalize()
# target line
tline = stimpoint - targetpoint
if tline.norm() < 5:
  tline = ntarget
  print('target is on the head: target line is the normal')
else:
  tline.normalize()
cosalpha = tline.dot( ntarget )
if cosalpha > 1.:
  cosalpha = 1.
alpha = math.acos( cosalpha )
print('stimulation position on head:', stimpoint.list())
print('normal:', ntarget.list())
print('target line:', tline.list())
print('(normal, target line) angle:', alpha)
print()
endpoint = stimpoint - ntarget * depth
if options.writeresults:
  f = open( options.writeresults, 'w' )
  print('target:', targetpoint.list(), file=f)
  print('coil:', stimpoint.list(), file=f)
  print('normal:', ntarget.list(), file=f)
  print('target_line:', tline.list(), file=f)
  print('normal_target_angle:', alpha, file=f)
  f.close()
  del f

stimarea = aims.SurfaceGenerator.cone( endpoint, stimpoint, baseradius, 50,
  False, True )

aims.write( stimarea, outfile )
if options.normalmesh:
  nmesh = aims.AimsTimeSurface_2()
  v0 = nmesh.vertex(0)
  v0.assign( [ endpoint, stimpoint + ntarget * 50 ] )
  p0 = nmesh.polygon(0)
  p0.append( [ 0, 1 ] )
  aims.write( nmesh, options.normalmesh )

if options.targetlinemesh:
  tlmesh = aims.AimsTimeSurface_2()
  v0 = tlmesh.vertex(0)
  tortho = aims.vectProduct( tline, [ 0., 0., 1. ] )
  if tortho.norm() < 1.e-4:
    tortho = aims.vectProduct( tline, [ 1., 0., 0. ] )
  tortho.normalize()
  tortho2 = aims.vectProduct( tline, tortho )
  tortho2.normalize() # (should be already normalized)
  v0.assign( [ aims.Point3df( targetpoint ) - tline * 50,
    targetpoint,
    stimpoint + tline * 50,
    targetpoint - tortho * 20,
    targetpoint + tortho * 20,
    targetpoint - tortho2 * 20,
    targetpoint + tortho2 * 20 ] )
  p0 = tlmesh.polygon(0)
  p0.append( [ 0, 1 ] )
  p0.append( [ 1, 2 ] )
  p0.append( [ 3, 4 ] )
  p0.append( [ 5, 6 ] )
  aims.write( tlmesh, options.targetlinemesh )

if options.stimzone:
  print('masking stimulation zone...')
  # read brain mask
  brain = aims.read( options.brain )
  vs = brain.header()[ 'voxel_size' ]
  dims = brain.header()[ 'volume_dimension' ]
  dimm = [ dims[i] * vs[i] for i in range( 3 ) ]
  # intersection with the stimulation zone
  #  (take a rough bounding box)
  bmin = numpy.min( ( stimpoint, endpoint ), axis=0 )
  bmax = numpy.max( ( stimpoint, endpoint ), axis=0 )
  bmin -= baseradius
  bmax += baseradius
  bmin = numpy.max( ( bmin, [ 0, 0, 0 ] ), axis=0 )
  bmax = numpy.min( ( bmax, dimm ), axis=0 )

  cs = aims.ConeSamplable( endpoint, stimpoint, baseradius )
  sampler = aimsalgo.BucketMapSampler_FLOAT_3()
  zonespl = sampler.sample( cs, bmin, bmax - bmin, vs[:3] ).getPython()
  if options.cortexvalue is None:
    stzone = aims.maskWithVolume( zonespl._get(), brain )
  else:
    stzone = aims.maskWithVolume( zonespl._get(), brain, options.cortexvalue,
      True )

  zone = stzone[0]
  print('stimulation zone:', zone.size(), 'voxels (',
        zone.size() * vs[0] * vs[1] * vs[2], 'mm3)')
  aims.write( stzone, options.stimzone )


# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Helper functions and functions for plotting gamma-ray images."""
import numpy as np
from astropy.coordinates import Angle

__all__ = ["colormap_hess", "colormap_milagro", "MapPanelPlotter"]

__doctest_requires__ = {("colormap_hess", "colormap_milagro"): ["matplotlib"]}

class MapPanelPlotter:
    """
    Map panel plotter class.

    Given a `~matplotlib.pyplot.Figure` object this class creates axes objects
    using `~matplotlib.gridspec.GridSpec` and plots a given sky map onto these.

    Parameters
    ----------
    figure : `~matplotlib.pyplot.figure.`
        Figure instance.
    xlim : `~astropy.coordinates.Angle`
        Angle object specifying the longitude limits.
    ylim : `~astropy.coordinates.Angle`
        Angle object specifying the latitude limits.
    npanels : int
        Number of panels.
    **kwargs : dict
        Keyword arguments passed to `~matplotlib.gridspec.GridSpec`.
    """

    def __init__(self, figure, xlim, ylim, npanels=4, **kwargs):
        from matplotlib.gridspec import GridSpec

        self.figure = figure
        self.parameters = {"xlim": xlim, "ylim": ylim, "npanels": npanels}
        self.grid_spec = GridSpec(nrows=npanels, ncols=1, **kwargs)

    def _get_ax_extend(self, ax, panel):
        """Get width and height of the axis in world coordinates."""
        p = self.parameters

        # compute aspect ratio of the axis
        aspect = ax.bbox.width / ax.bbox.height

        # compute width and height in world coordinates
        height = np.abs(p["ylim"].diff())
        width = aspect * height

        left, bottom = p["xlim"][0].wrap_at("180d"), p["ylim"][0]

        width_all = np.abs(p["xlim"].wrap_at("180d").diff())
        xoverlap = ((p["npanels"] * width) - width_all) / (p["npanels"] - 1.0)
        if xoverlap < 0:
            raise ValueError(
                "No overlap between panels. Please reduce figure "
                "height or increase vertical space between the panels."
            )

        left = left - panel * (width - xoverlap)
        return left[0], bottom, width, height

    def _set_ax_fov(self, ax, panel):
        left, bottom, width, height = self._get_ax_extend(ax, panel)

        # set fov
        xlim = Angle([left, left - width])
        ylim = Angle([bottom, bottom + height])
        xlim_pix, ylim_pix = ax.wcs.wcs_world2pix(xlim.deg, ylim.deg, 1)

        ax.set_xlim(*xlim_pix)
        ax.set_ylim(*ylim_pix)
        return ax

    def plot_panel(self, map, panel=1, panel_fov=None, **kwargs):
        """
        Plot sky map on one panel.

        Parameters
        ----------
        map : `~gammapy.maps.WcsNDMap`
            Map to plot.
        panel : int
            Which panel to plot on (counted from top).
        """
        if panel_fov is None:
            panel_fov = panel
        spec = self.grid_spec[panel]
        ax = self.figure.add_subplot(spec, projection=map.geom.wcs)
        try:
            ax = map.plot(ax=ax, **kwargs)[1]
        except AttributeError:
            ax = map.plot_rgb(ax=ax, **kwargs)
        ax = self._set_ax_fov(ax, panel_fov)
        return ax

    def plot(self, map, **kwargs):
        """
        Plot sky map on all panels.

        Parameters
        ----------
        map : `~gammapy.maps.WcsNDMap`
            Map to plot.
        """
        p = self.parameters
        axes = []
        for panel in range(p["npanels"]):
            ax = self.plot_panel(map, panel=panel, **kwargs)
            axes.append(ax)
        return axes

def colormap_hess(transition=0.5, width=0.1):
    """Colormap often used in H.E.S.S. collaboration publications.

    This colormap goes black -> blue -> red -> yellow -> white.

    A sharp blue -> red -> yellow transition is often used for significance images
    with a value of red at ``transition ~ 5`` or ``transition ~ 7``
    so that the following effect is achieved:

    - black, blue: non-significant features, not well visible
    - red: features at the detection threshold ``transition``
    - yellow, white: significant features, very well visible

    The transition parameter is defined between 0 and 1. To calculate the value
    from data units an `~astropy.visualization.mpl_normalize.ImageNormalize`
    instance should be used (see example below).

    Parameters
    ----------
    transition : float (default = 0.5)
        Value of the transition to red (between 0 and 1).
    width : float (default = 0.5)
        Width of the blue-red color transition (between 0 and 1).

    Returns
    -------
    colormap : `matplotlib.colors.LinearSegmentedColormap`
        Colormap

    Examples
    --------
    >>> from gammapy.maps import colormap_hess
    >>> from astropy.visualization.mpl_normalize import ImageNormalize
    >>> from astropy.visualization import LinearStretch
    >>> normalize = ImageNormalize(vmin=-5, vmax=15, stretch=LinearStretch())
    >>> transition = normalize(5)
    >>> cmap = colormap_hess(transition=transition)
    """
    from matplotlib.colors import LinearSegmentedColormap

    # Compute normalised values (range 0 to 1) that
    # correspond to red, blue, yellow.
    red = float(transition)

    if width > red:
        blue = 0.1 * red
    else:
        blue = red - width

    yellow = 2.0 / 3.0 * (1 - red) + red

    black, white = 0, 1

    # Create custom colormap
    # List entries: (value, (R, G, B))
    colors = [
        (black, "k"),
        (blue, (0, 0, 0.8)),
        (red, "r"),
        (yellow, (1.0, 1.0, 0)),
        (white, "w"),
    ]

    return LinearSegmentedColormap.from_list(name="hess", colors=colors)

def colormap_milagro(transition=0.5, width=0.0001, huestart=0.6):
    """Colormap often used in Milagro collaboration publications.

    This colormap is gray below ``transition`` and similar to the jet colormap above.

    A sharp gray -> color transition is often used for significance images
    with a transition value of ``transition ~ 5`` or ``transition ~ 7``,
    so that the following effect is achieved:

    - gray: non-significant features are not well visible
    - color: significant features at the detection threshold ``transition``

    Note that this colormap is often criticised for over-exaggerating small differences
    in significance below and above the gray - color transition threshold.

    The transition parameter is defined between 0 and 1. To calculate the value
    from data units an `~astropy.visualization.mpl_normalize.ImageNormalize` instance should be
    used (see example below).

    Parameters
    ----------
    transition : float (default = 0.5)
        Transition value (below: gray, above: color).
    width : float (default = 0.0001)
        Width of the transition
    huestart : float (default = 0.6)
        Hue of the color at ``transition``

    Returns
    -------
    colormap : `~matplotlib.colors.LinearSegmentedColormap`
        Colormap

    Examples
    --------
    >>> from gammapy.maps import colormap_milagro
    >>> from astropy.visualization.mpl_normalize import ImageNormalize
    >>> from astropy.visualization import LinearStretch
    >>> normalize = ImageNormalize(vmin=-5, vmax=15, stretch=LinearStretch())
    >>> transition = normalize(5)
    >>> cmap = colormap_milagro(transition=transition)
    """
    from colorsys import hls_to_rgb
    from matplotlib.colors import LinearSegmentedColormap

    # Compute normalised red, blue, yellow values
    transition = float(transition)

    # Create custom colormap
    # List entries: (value, (H, L, S))
    colors = [
        (0, (1, 1, 0)),
        (transition - width, (1, 0, 0)),
        (transition, (huestart, 0.4, 0.5)),
        (transition + width, (huestart, 0.4, 1)),
        (0.99, (0, 0.6, 1)),
        (1, (0, 1, 1)),
    ]

    # Convert HLS values to RGB values
    rgb_colors = [(val, hls_to_rgb(*hls)) for (val, hls) in colors]

    return LinearSegmentedColormap.from_list(name="milagro", colors=rgb_colors)


import os

import cv2
import albumentations as alb
import torch
import torchvision.datasets as cv_datasets
import torchvision.ops as cv_ops
import pycocotools.mask as coco_mask
import numpy as np

import datasets.tools as tools
import utils.bbox_ops as bbox_ops

class TrainDataset(cv_datasets.CocoDetection):
    def __init__(self, root, year, input_size, pooler_size):
        super(TrainDataset, self).__init__(root=os.path.join(root, f'train{year}'), annFile=os.path.join(root, 'annotations', f'instances_train{year}.json'))

        self.h, self.w = input_size
        self.pooler_size = pooler_size

        self.cat_idx_list, self.cat_to_label_map, _ = tools.get_cat_label_map(self.coco, tools.COCO_CLASSES)

        self.img_transform = alb.Compose([
            alb.RandomSizedBBoxSafeCrop(width=self.w, height=self.h),
            alb.HorizontalFlip(p=0.5),
            alb.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0., p=0.8),
        ], bbox_params=alb.BboxParams(format='coco', label_fields=['class_labels']))

        self.points, self.regress_ranges = tools.encode_points_and_regress_ranges(self.h, self.w)

    def __getitem__(self, index):
        img, target = tools.load_img_target(self, index)
        img_info = self.coco.loadImgs(self.ids[index])[0]
        iw, ih = img_info['width'], img_info['height']

        class_labels, bbox_labels, mask_labels = [], [], []
        for obj in target:
            if not tools.is_correct_instance(obj, self.cat_idx_list, iw, ih):
                continue

            class_labels.append(self.cat_to_label_map[obj['category_id']])
            bbox_labels.append(obj['bbox'])

            # rle = coco_mask.frPyObjects(obj['segmentation'], ih, iw)
            # if obj['iscrowd'] == 0:
            #     rle = coco_mask.merge(rle)
            # mask = coco_mask.decode(rle)
            # mask_labels.append(mask)

        transformed = self.img_transform(image=img, bboxes=bbox_labels, class_labels=class_labels)
        # transformed = self.img_transform(image=img, masks=mask_labels, bboxes=bbox_labels, class_labels=class_labels)
        img = tools.TENSOR_TRANSFORM(transformed['image'])
        # mask_labels = transformed['masks']
        class_labels = transformed['class_labels']
        bbox_labels = transformed['bboxes']

        if len(bbox_labels) == 0:
            # For any instance with classification label 0 (background), only classification loss will be computed, without mask loss, centerness loss and bbox loss.
            # When there is no instances in an image, it doesn't matter the value of the added bbox.
            mask_labels = [np.zeros((self.h, self.w))]
            bbox_labels = [[0., 0., 10., 10.]]
            class_labels = [0]

        class_labels = torch.as_tensor(class_labels)

        # instance_mask_labels = self._generate_instance_mask_labels(mask_labels, bbox_labels)
        # instance_mask_labels = torch.as_tensor(np.array(instance_mask_labels)).float()

        bbox_labels = cv_ops.box_convert(torch.as_tensor(bbox_labels, dtype=torch.float32), in_fmt='xywh', out_fmt='xyxy')
        bbox_labels = cv_ops.clip_boxes_to_image(bbox_labels, (ih, iw))

        class_targets, distance_targets = self._encode_targets(class_labels, bbox_labels, None)
        centerness_targets = tools.encode_centerness_targets(distance_targets)

        return img, self.points, {'class': class_targets, 'distance': distance_targets, 'centerness': centerness_targets}

    def _generate_instance_mask_labels(self, mask_labels, bbox_labels):
        instance_mask_labels = []
        for mask_label, bbox_label in zip(mask_labels, bbox_labels):
            bx, by, bw, bh = bbox_label
            bx, by, bw, bh = int(bx), int(by), int(bw), int(bh)

            roi_mask = mask_label[by:by + bh, bx:bx + bw]
            roi_mask = cv2.resize(roi_mask, dsize=(self.pooler_size, self.pooler_size))
            instance_mask_labels.append(roi_mask)
        return instance_mask_labels

    def _encode_targets(self, cls_labels, bbox_labels, instance_mask_labels):
        points = self.points.clone()
        regress_ranges = self.regress_ranges.clone()

        num_points = points.size(0)
        num_gts = cls_labels.size(0)

        regress_ranges = regress_ranges[:, None, :].repeat(1, num_gts, 1)  # [num_points, num_gts, 2]
        bbox_areas = cv_ops.box_area(bbox_labels)[None].repeat(num_points, 1)  # [num_points, num_gts]

        expanded_points = points[:, None, :].repeat(1, num_gts, 1)
        expanded_bboxes = bbox_labels[None, :, :].repeat(num_points, 1, 1)
        distance_targets = bbox_ops.convert_bbox_to_distance(expanded_points, expanded_bboxes)  # [num_points, num_gts, 4]
        # instance_mask_labels = instance_mask_labels[None, :, :, :].repeat(num_points, 1, 1, 1)  # [num_points, num_gts, roi_size, roi_size]

        # Condition 1: inside a gt bbox
        inside_gt_bbox_mask = distance_targets.min(dim=-1)[0] > 0  # [num_points, num_gts]

        # Condition 2: limit the regression range for each location
        max_regress_distance = distance_targets.max(dim=-1)[0]  # [num_points, num_gts]
        inside_regress_range = (max_regress_distance >= regress_ranges[..., 0]) & (max_regress_distance <= regress_ranges[..., 1])  # [num_points, num_gts]

        # If there are still more than one instances for a location, we choose the one with minimal area
        bbox_areas[inside_gt_bbox_mask == 0] = tools.INF
        bbox_areas[inside_regress_range == 0] = tools.INF
        min_area, min_area_idx = bbox_areas.min(dim=1)  # [num_points], Assign a gt to each location

        class_targets = cls_labels[min_area_idx]
        class_targets[min_area == tools.INF] = 0

        distance_targets = distance_targets[range(num_points), min_area_idx, :]
        # instance_mask_labels = instance_mask_labels[range(num_points), min_area_idx, :, :]

        return class_targets, distance_targets  # , instance_mask_labels