import os
import time
from pathlib import Path

import numpy as np

np.warnings.filterwarnings('ignore')
import flopy
import pandas as pd

fm = flopy.modflow
from flopy.modflow import Modflow

from mfsetup.bcs import (
    remove_inactive_bcs,
    setup_flopy_stress_period_data,
    setup_ghb_data,
)
from mfsetup.discretization import (
    deactivate_idomain_above,
    find_remove_isolated_cells,
    make_ibound,
)
from mfsetup.fileio import (
    add_version_to_fileheader,
    flopy_mf2005_load,
    load,
    load_cfg,
    save_array,
)
from mfsetup.ic import setup_strt
from mfsetup.lakes import (
    make_bdlknc2d,
    make_bdlknc_zones,
    setup_lake_fluxes,
    setup_lake_info,
    setup_lake_tablefiles,
)
from mfsetup.mfmodel import MFsetupMixin
from mfsetup.obs import read_observation_data, setup_head_observations
from mfsetup.oc import parse_oc_period_input
from mfsetup.tdis import get_parent_stress_periods, setup_perioddata_group
from mfsetup.tmr import TmrNew
from mfsetup.units import convert_length_units, itmuni_text, lenuni_text
from mfsetup.utils import get_input_arguments, get_packages

from .wells import setup_wel_data

class MFnwtModel(MFsetupMixin, Modflow):
    """Class representing a MODFLOW-NWT model"""
    default_file = '/mfnwt_defaults.yml'

    def __init__(self, parent=None, cfg=None,
                 modelname='model', exe_name='mfnwt',
                 version='mfnwt', model_ws='.',
                 external_path='external/', **kwargs):
        defaults = {'parent': parent,
                    'modelname': modelname,
                    'exe_name': exe_name,
                    'version': version,
                    'model_ws': model_ws,
                    'external_path': external_path,
                    }
        # load configuration, if supplied
        if cfg is not None:
            if not isinstance(cfg, dict):
                cfg = self.load_cfg(cfg)
            cfg = self._parse_model_kwargs(cfg)
            defaults.update(cfg['model'])
            kwargs = {k: v for k, v in kwargs.items() if k not in defaults}
        # otherwise, pass arguments on to flopy constructor
        args = get_input_arguments(defaults, Modflow,
                                   exclude='packages')
        Modflow.__init__(self, **args, **kwargs)
        #Modflow.__init__(self, modelname, exe_name=exe_name, version=version,
        #                 model_ws=model_ws, external_path=external_path,
        #                 **kwargs)
        MFsetupMixin.__init__(self, parent=parent)

        # default configuration
        self._package_setup_order = ['dis', 'bas6', 'upw', 'rch', 'oc',
                                     'chd', 'ghb', 'lak', 'sfr', 'riv', 'wel', 'mnw2',
                                     'gag', 'hyd']
        # set up the model configuration dictionary
        # start with the defaults
        self.cfg = load(self.source_path + self.default_file)  # '/mf6_defaults.yml')
        self.relative_external_paths = self.cfg.get('model', {}).get('relative_external_paths', True)
        # set the model workspace and change working directory to there
        self.model_ws = self._get_model_ws(cfg=cfg)
        # update defaults with user-specified config. (loaded above)
        # set up and validate the model configuration dictionary
        self._set_cfg(cfg)

        # set the list file path
        self.lst.file_name = [self.cfg['model']['list_filename_fmt'].format(self.name)]

        # property arrays
        self._ibound = None

    def __repr__(self):
        return MFsetupMixin.__repr__(self)

    @property
    def nlay(self):
        return self.cfg['dis'].get('nlay', 1)

    @property
    def length_units(self):
        return lenuni_text[self.cfg['dis']['lenuni']]

    @property
    def time_units(self):
        return itmuni_text[self.cfg['dis']['itmuni']]

    @property
    def ipakcb(self):
        """By default write everything to one cell budget file."""
        return self.cfg['upw'].get('ipakcb', 53)

    @property
    def ibound(self):
        """3D array indicating which cells will be included in the simulation.
        Made a property so that it can be easily updated when any packages
        it depends on change.
        """
        if self._ibound is None and 'BAS6' in self.get_package_list():
            self._set_ibound()
        return self._ibound

    def _set_ibound(self):
        """Remake the idomain array from the source data,
        no data values in the top and bottom arrays, and
        so that cells above SFR reaches are inactive."""
        ibound_from_layer_elevations = make_ibound(self.dis.top.array,
                                                     self.dis.botm.array,
                                                     nodata=self._nodata_value,
                                                     minimum_layer_thickness=self.cfg['dis'].get(
                                                         'minimum_layer_thickness', 1),
                                                     #drop_thin_cells=self._drop_thin_cells,
                                                     tol=1e-4)

        # include cells that are active in the existing idomain array
        # and cells inactivated on the basis of layer elevations
        ibound = (self.bas6.ibound.array > 0) & (ibound_from_layer_elevations >= 1)
        ibound = ibound.astype(int)

        # remove cells that conincide with lakes
        ibound[self.isbc == 1] = 0.

        # remove cells that are above stream cells
        if self.get_package('sfr') is not None:
            ibound = deactivate_idomain_above(ibound, self.sfr.reach_data)
        # remove cells that are above ghb cells
        if self.get_package('ghb') is not None:
            ibound = deactivate_idomain_above(ibound, self.ghb.stress_period_data[0])

        # inactivate any isolated cells that could cause problems with the solution
        ibound = find_remove_isolated_cells(ibound, minimum_cluster_size=20)

        self._ibound = ibound
        # re-write the input files
        self._setup_array('bas6', 'ibound', resample_method='nearest',
                          data={i: arr for i, arr in enumerate(ibound)},
                          datatype='array3d', write_fmt='%d', dtype=int)
        self.bas6.ibound = self.cfg['bas6']['ibound']

    def _set_parent(self):
        """Set attributes related to a parent or source model
        if one is specified."""

        if self.cfg['parent'].get('version') == 'mf6':
            raise NotImplementedError("MODFLOW-6 parent models")

        kwargs = self.cfg['parent'].copy()
        if kwargs is not None:
            kwargs = kwargs.copy()
            kwargs['f'] = kwargs.pop('namefile')
            # load only specified packages that the parent model has
            packages_in_parent_namefile = get_packages(os.path.join(kwargs['model_ws'],
                                                                    kwargs['f']))
            load_only = list(set(packages_in_parent_namefile).intersection(
                set(self.cfg['model'].get('packages', set()))))
            kwargs['load_only'] = load_only
            kwargs = get_input_arguments(kwargs, fm.Modflow.load, warn=False)

            print('loading parent model {}...'.format(os.path.join(kwargs['model_ws'],
                                                                   kwargs['f'])))
            t0 = time.time()
            self._parent = fm.Modflow.load(**kwargs)
            print("finished in {:.2f}s\n".format(time.time() - t0))

            # parent model units
            if 'length_units' not in self.cfg['parent']:
                self.cfg['parent']['length_units'] = lenuni_text[self.parent.dis.lenuni]
            if 'time_units' not in self.cfg['parent']:
                self.cfg['parent']['time_units'] = itmuni_text[self.parent.dis.itmuni]

            # set the parent model grid from mg_kwargs if not None
            # otherwise, convert parent model grid to MFsetupGrid
            mg_kwargs = self.cfg['parent'].get('SpatialReference',
                                          self.cfg['parent'].get('modelgrid', None))
            self._set_parent_modelgrid(mg_kwargs)

            # parent model perioddata
            if not hasattr(self.parent, 'perioddata'):
                kwargs = {}
                kwargs['start_date_time'] = self.cfg['parent'].get('start_date_time',
                                                                   self.cfg['model'].get('start_date_time',
                                                                                '1970-01-01'))
                kwargs['nper'] = self.parent.nper
                kwargs['model_time_units'] = self.cfg['parent']['time_units']
                for var in ['perlen', 'steady', 'nstp', 'tsmult']:
                    kwargs[var] = self.parent.dis.__dict__[var].array
                kwargs = get_input_arguments(kwargs, setup_perioddata_group)
                kwargs['oc_saverecord'] = {}
                self._parent.perioddata = setup_perioddata_group(**kwargs)

            # default_source_data, where omitted configuration input is
            # obtained from parent model by default
            # Set default_source_data to True by default if it isn't specified
            if self.cfg['parent'].get('default_source_data') is None:
                self.cfg['parent']['default_source_data'] = True
            if self.cfg['parent'].get('default_source_data'):
                self._parent_default_source_data = True
                if self.cfg['dis'].get('nlay') is None:
                    self.cfg['dis']['nlay'] = self.parent.dis.nlay
                parent_start_date_time = self.cfg.get('parent', {}).get('start_date_time')
                if self.cfg['dis'].get('start_date_time', '1970-01-01') == '1970-01-01' and parent_start_date_time is not None:
                    self.cfg['dis']['start_date_time'] = self.cfg['parent']['start_date_time']
                if self.cfg['dis'].get('nper') is None:
                    self.cfg['dis']['nper'] = self.parent.dis.nper
                parent_periods = get_parent_stress_periods(self.parent, nper=self.cfg['dis']['nper'],
                                                           parent_stress_periods=self.cfg['parent']['copy_stress_periods'])
                for var in ['perlen', 'nstp', 'tsmult', 'steady']:
                    if self.cfg['dis'].get(var) is None:
                        self.cfg['dis'][var] = self.parent.dis.__dict__[var].array[parent_periods]

    def _update_grid_configuration_with_dis(self):
        """Update grid configuration with any information supplied to dis package
        (so that settings specified for DIS package have priority). This method
        is called by MFsetupMixin.setup_grid.
        """
        for param in ['nrow', 'ncol', 'delr', 'delc']:
            if param in self.cfg['dis']:
                self.cfg['setup_grid'][param] = self.cfg['dis'][param]

    def setup_dis(self):
        """"""
        package = 'dis'
        print('\nSetting up {} package...'.format(package.upper()))
        t0 = time.time()

        # resample the top from the DEM
        if self.cfg['dis']['remake_top']:
            self._setup_array(package, 'top', datatype='array2d',
                              resample_method='linear',
                              write_fmt='%.2f')

        # make the botm array
        self._setup_array(package, 'botm', datatype='array3d',
                          resample_method='linear',
                          write_fmt='%.2f')

        # put together keyword arguments for dis package
        kwargs = self.cfg['grid'].copy() # nrow, ncol, delr, delc
        kwargs.update(self.cfg['dis']) # nper, nlay, etc.
        kwargs = get_input_arguments(kwargs, fm.ModflowDis)
        # we need flopy to read the intermediate files
        # (it will write the files in cfg)
        lmult = convert_length_units('meters', self.length_units)
        kwargs.update({'top': self.cfg['intermediate_data']['top'][0],
                       'botm': self.cfg['intermediate_data']['botm'],
                       'nper': self.nper,
                       'delc': self.modelgrid.delc * lmult,
                       'delr': self.modelgrid.delr * lmult
                      })
        for arg in ['perlen', 'nstp', 'tsmult', 'steady']:
            kwargs[arg] = self.perioddata[arg].values

        dis = fm.ModflowDis(model=self, **kwargs)
        self._perioddata = None  # reset perioddata
        #if not isinstance(self._modelgrid, MFsetupGrid):
        #    self._modelgrid = None  # override DIS package grid setup
        self.setup_grid()  # reset the model grid
        self._reset_bc_arrays()
        #self._isbc = None  # reset BC property arrays
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return dis

    def setup_tdis(self):
        """Calls the _set_perioddata, to establish time discretization. Only purpose
        is to conform to same syntax as mf6 for MFsetupMixin.setup_from_yaml()
        """
        self._set_perioddata()

    def setup_bas6(self):
        """"""
        package = 'bas6'
        print('\nSetting up {} package...'.format(package.upper()))
        t0 = time.time()

        kwargs = self.cfg[package]
        kwargs['source_data_config'] = kwargs['source_data']
        kwargs['filename_fmt'] = kwargs['strt_filename_fmt']
        kwargs['write_fmt'] = kwargs['strt_write_fmt']

        # make the starting heads array
        strt = setup_strt(self, package, **kwargs)

        # initial ibound input for creating a bas6 package instance
        self._setup_array(package, 'ibound', datatype='array3d', write_fmt='%d',
                          resample_method='nearest',
                          dtype=int)

        kwargs = get_input_arguments(self.cfg['bas6'], fm.ModflowBas)
        kwargs['strt'] = strt
        bas = fm.ModflowBas(model=self, **kwargs)
        print("finished in {:.2f}s\n".format(time.time() - t0))
        self._set_ibound()
        return bas

    def setup_oc(self):

        package = 'oc'
        print('\nSetting up {} package...'.format(package.upper()))
        t0 = time.time()
        #stress_period_data = {}
        #for i, r in self.perioddata.iterrows():
        #    stress_period_data[(r.per, r.nstp -1)] = r.oc

        # use stress_period_data if supplied
        # (instead of period_input defaults)
        if 'stress_period_data' in self.cfg['oc']:
            del self.cfg['oc']['period_options']
        kwargs = self.cfg['oc']
        period_input = parse_oc_period_input(kwargs, nstp=self.perioddata.nstp,
                                             output_fmt='mfnwt')
        kwargs.update(period_input)
        kwargs = get_input_arguments(kwargs, fm.ModflowOc)
        oc = fm.ModflowOc(model=self, **kwargs)
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return oc

    def setup_rch(self):
        package = 'rch'
        print('\nSetting up {} package...'.format(package.upper()))
        t0 = time.time()

        # make the rech array
        self._setup_array(package, 'rech', datatype='transient2d',
                          resample_method='linear',
                          write_fmt='%.6e',
                          write_nodata=0.)

        # create flopy package instance
        kwargs = self.cfg['rch']
        kwargs['ipakcb'] = self.ipakcb
        kwargs = get_input_arguments(kwargs, fm.ModflowRch)
        rch = fm.ModflowRch(model=self, **kwargs)
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return rch

    def setup_upw(self):
        """
        """
        package = 'upw'
        print('\nSetting up {} package...'.format(package.upper()))
        t0 = time.time()
        hiKlakes_value = float(self.cfg['parent'].get('hiKlakes_value', 1e4))

        # copy transient variables if they were included in config file
        # defaults are hard coded to arrays in parent model priority
        # over config file values, in the case that ss and sy weren't entered
        hk = self.cfg['upw'].get('hk')
        vka = self.cfg['upw'].get('vka')
        default_sy = 0.1
        default_ss = 1e-6

        # Determine which hk, vka to use
        # load parent upw if it's needed and not loaded
        source_package = package
        if np.any(np.array([hk, vka]) == None) and \
                'UPW' not in self.parent.get_package_list() and \
                'LPF' not in self.parent.get_package_list():
            for ext, pckgcls in {'upw': fm.ModflowUpw,
                                 'lpf': fm.ModflowLpf,
                                 }.items():
                pckgfile = '{}/{}.{}'.format(self.parent.model_ws, self.parent.name, package)
                if os.path.exists(pckgfile):
                    upw = pckgcls.load(pckgfile, self.parent)
                    source_package = ext
                    break

        self._setup_array(package, 'hk', vmin=0, vmax=hiKlakes_value, resample_method='linear',
                           source_package=source_package, datatype='array3d', write_fmt='%.6e')
        self._setup_array(package, 'vka', vmin=0, vmax=hiKlakes_value, resample_method='linear',
                           source_package=source_package, datatype='array3d', write_fmt='%.6e')
        if np.any(~self.dis.steady.array):
            self._setup_array(package, 'sy', vmin=0, vmax=1, resample_method='linear',
                              source_package=source_package,
                              datatype='array3d', write_fmt='%.6e')
            self._setup_array(package, 'ss', vmin=0, vmax=1, resample_method='linear',
                              source_package=source_package,
                              datatype='array3d', write_fmt='%.6e')
            sy = self.cfg['intermediate_data']['sy']
            ss = self.cfg['intermediate_data']['ss']
        else:
            sy = default_sy
            ss = default_ss

        upw = fm.ModflowUpw(self, hk=self.cfg['intermediate_data']['hk'],
                            vka=self.cfg['intermediate_data']['vka'],
                            sy=sy,
                            ss=ss,
                            layvka=self.cfg['upw']['layvka'],
                            laytyp=self.cfg['upw']['laytyp'],
                            hdry=self.cfg['upw']['hdry'],
                            ipakcb=self.cfg['upw']['ipakcb'])
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return upw

    def setup_wel(self):
        """
        Setup the WEL package, including boundary fluxes and any pumping.

        This will need some additional customization if pumping is mixed with
        the perimeter fluxes.

        TODO: generalize well package setup with specific input requirements

        """

        print('setting up WEL package...')
        t0 = time.time()

        # munge the well package input
        # for_external_files only needs to be called on the modflow-6 side
        df = setup_wel_data(self, for_external_files=False)

        # extend spd dtype to include comments
        dtype = fm.ModflowWel.get_default_dtype()

        # setup stress period data
        groups = df.groupby('per')
        spd = {}
        for per, perdf in groups:
            ra = np.recarray(len(perdf), dtype=dtype)
            for c in ['k', 'i', 'j']:
                ra[c] = perdf[c]
            ra['flux'] = perdf['q']
            spd[per] = ra

        wel = fm.ModflowWel(self, ipakcb=self.ipakcb,
                            options=self.cfg['wel']['options'],
                            stress_period_data=spd)
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return wel

    def setup_ghb(self):
        """
        Set up the GHB package
        """

        print('setting up GHB package...')
        t0 = time.time()

        df = setup_ghb_data(self)

        # extend spd dtype to include comments
        dtype = fm.ModflowGhb.get_default_dtype()

        # setup stress period data
        groups = df.groupby('per')
        spd = {}
        for per, perdf in groups:
            ra = np.recarray(len(perdf), dtype=dtype)
            for c in ['k', 'i', 'j', 'bhead', 'cond']:
                ra[c] = perdf[c]
            spd[per] = ra

        ghb = fm.ModflowGhb(self, ipakcb=self.ipakcb,
                            stress_period_data=spd)
        self._reset_bc_arrays()
        self._ibound = None
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return ghb

    def setup_mnw2(self):

        print('setting up MNW2 package...')
        t0 = time.time()

        # added wells
        # todo: generalize MNW2 source data input; add auto-reprojection
        added_wells = self.cfg['mnw'].get('added_wells')
        if added_wells is not None:
            if isinstance(added_wells, str):
                aw = pd.read_csv(added_wells)
                aw.rename(columns={'name': 'comments'}, inplace=True)
            elif isinstance(added_wells, dict):
                added_wells = {k: v for k, v in added_wells.items() if v is not None}
                if len(added_wells) > 0:
                    aw = pd.DataFrame(added_wells).T
                    aw['comments'] = aw.index
                else:
                    aw = None
            elif isinstance(added_wells, pd.DataFrame):
                aw = added_wells
                aw['comments'] = aw.index
            else:
                raise IOError('unrecognized added_wells input')

            k, ztop, zbotm = 0, 0, 0
            zpump = None

            wells = aw.groupby('comments').first()
            periods = aw
            if 'x' in wells.columns and 'y' in wells.columns:
                wells['i'], wells['j'] = self.modelgrid.intersect(wells['x'].values,
                                                     wells['y'].values)
            if 'depth' in wells.columns:
                wellhead_elevations = self.dis.top.array[wells.i, wells.j]
                ztop = wellhead_elevations - (5*.3048) # 5 ft casing
                zbotm = wellhead_elevations - wells.depth
                zpump = zbotm + 1 # 1 meter off bottom
            elif 'ztop' in wells.columns and 'zbotm' in wells.columns:
                ztop = wells.ztop
                zbotm = wells.zbotm
                zpump = zbotm + 1
            if 'k' in wells.columns:
                k = wells.k

            for var in ['losstype', 'pumploc', 'rw', 'rskin', 'kskin']:
                if var not in wells.columns:
                    wells[var] = self.cfg['mnw']['defaults'][var]

            nd = fm.ModflowMnw2.get_empty_node_data(len(wells))
            nd['k'] = k
            nd['i'] = wells.i
            nd['j'] = wells.j
            nd['ztop'] = ztop
            nd['zbotm'] = zbotm
            nd['wellid'] = wells.index
            nd['losstype'] = wells.losstype
            nd['pumploc'] = wells.pumploc
            nd['rw'] = wells.rw
            nd['rskin'] = wells.rskin
            nd['kskin'] = wells.kskin
            if zpump is not None:
                nd['zpump'] = zpump

            spd = {}
            for per, group in periods.groupby('per'):
                spd_per = fm.ModflowMnw2.get_empty_stress_period_data(len(group))
                spd_per['wellid'] = group.comments
                spd_per['qdes'] = group.flux
                spd[per] = spd_per
            itmp = []
            for per in range(self.nper):
                if per in spd.keys():
                    itmp.append(len(spd[per]))
                else:
                    itmp.append(0)

            mnw = fm.ModflowMnw2(self, mnwmax=len(wells), ipakcb=self.ipakcb,
                                 mnwprnt=1,
                                 node_data=nd, stress_period_data=spd,
                                 itmp=itmp
                                 )
            print("finished in {:.2f}s\n".format(time.time() - t0))
            return mnw
        else:
            print('No wells specified in configuration file!\n')
            return None

    def setup_lak(self):

        print('setting up LAKE package...')
        t0 = time.time()
        # if shapefile of lakes was included,
        # lakarr should be automatically built by property method
        if self.lakarr.sum() == 0:
            print("lakes_shapefile not specified, or no lakes in model area")
            return

        # source data
        source_data = self.cfg['lak']['source_data']
        self.lake_info = setup_lake_info(self)
        nlakes = len(self.lake_info)

        # set up the tab files, if any
        tab_files_argument = None
        tab_units = None
        start_tab_units_at = 150  # default starting number for iunittab
        if 'stage_area_volume_file' in source_data:

            tab_files = setup_lake_tablefiles(self, source_data['stage_area_volume_file'])
            tab_units = list(range(start_tab_units_at, start_tab_units_at + len(tab_files)))

            # tabfiles aren't rewritten by flopy on package write
            self.cfg['lak']['tab_files'] = tab_files
            # kludge to deal with ugliness of lake package external file handling
            # (need to give path relative to model_ws, not folder that flopy is working in)
            tab_files_argument = [os.path.relpath(f) for f in tab_files]

        self.setup_external_filepaths('lak', 'lakzones',
                                      self.cfg['lak']['{}_filename_fmt'.format('lakzones')])
        self.setup_external_filepaths('lak', 'bdlknc',
                                      self.cfg['lak']['{}_filename_fmt'.format('bdlknc')],
                                      file_numbers=list(range(self.nlay)))

        # make the arrays or load them
        lakzones = make_bdlknc_zones(self.modelgrid, self.lake_info,
                                     include_ids=self.lake_info['feat_id'])
        save_array(self.cfg['intermediate_data']['lakzones'][0], lakzones, fmt='%d')

        bdlknc = np.zeros((self.nlay, self.nrow, self.ncol))
        # make the areal footprint of lakebed leakance from the zones (layer 1)
        bdlknc[0] = make_bdlknc2d(lakzones,
                                  self.cfg['lak']['source_data']['littoral_leakance'],
                                  self.cfg['lak']['source_data']['profundal_leakance'])
        for k in range(self.nlay):
            if k > 0:
                # for each underlying layer, assign profundal leakance to cells were isbc == 1
                bdlknc[k][self.isbc[k] == 1] = self.cfg['lak']['source_data']['profundal_leakance']
            save_array(self.cfg['intermediate_data']['bdlknc'][0][k], bdlknc[k], fmt='%.6e')

        # get estimates of stage from model top, for specifying ranges
        stages = []
        for lakid in self.lake_info['lak_id']:
            loc = self.lakarr[0] == lakid
            est_stage = self.dis.top.array[loc].min()
            stages.append(est_stage)
        stages = np.array(stages)

        # setup stress period data
        tol = 5  # specify lake stage range as +/- this value
        ssmn, ssmx = stages - tol, stages + tol
        stage_range = list(zip(ssmn, ssmx))

        # set up dataset 9
        # ssmn and ssmx values only required for steady-state periods > 0
        self.lake_fluxes = setup_lake_fluxes(self)
        precip = self.lake_fluxes['precipitation'].tolist()
        evap = self.lake_fluxes['evaporation'].tolist()
        flux_data = {}
        for i, steady in enumerate(self.dis.steady.array):
            if i > 0 and steady:
                flux_data_i = []
                for lake_ssmn, lake_ssmx in zip(ssmn, ssmx):
                    flux_data_i.append([precip[i], evap[i], 0, 0, lake_ssmn, lake_ssmx])
            else:
                flux_data_i = [[precip[i], evap[i], 0, 0]] * nlakes
            flux_data[i] = flux_data_i
        options = ['tableinput'] if tab_files_argument is not None else None

        kwargs = self.cfg['lak']
        kwargs['nlakes'] = len(self.lake_info)
        kwargs['stages'] = stages
        kwargs['stage_range'] = stage_range
        kwargs['flux_data'] = flux_data
        kwargs['tab_files'] = tab_files_argument  #This needs to be in the order of the lake IDs!
        kwargs['tab_units'] = tab_units
        kwargs['options'] = options
        kwargs['ipakcb'] = self.ipakcb
        kwargs['lwrt'] = 0
        kwargs = get_input_arguments(kwargs, fm.mflak.ModflowLak)
        lak = fm.ModflowLak(self, **kwargs)
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return lak

    def setup_nwt(self):

        print('setting up NWT package...')
        t0 = time.time()
        use_existing_file = self.cfg['nwt'].get('use_existing_file')
        kwargs = self.cfg['nwt']
        if use_existing_file is not None:
            #set use_existing_file relative to source path
            filepath = os.path.join(self._config_path,
                                    use_existing_file)

            assert os.path.exists(filepath), "Couldn't find {}, need a path to a NWT file".format(filepath)
            nwt = fm.ModflowNwt.load(filepath, model=self)
        else:
            kwargs = get_input_arguments(kwargs, fm.ModflowNwt)
            nwt = fm.ModflowNwt(self, **kwargs)
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return nwt

    def setup_hyd(self):
        """TODO: generalize hydmod setup with specific input requirements"""
        package = 'hyd'
        print('setting up HYDMOD package...')
        t0 = time.time()

        # munge the head observation data
        df = setup_head_observations(self, format=package,
                                     obsname_column='hydlbl')

        # create observation data recarray
        obsdata = fm.ModflowHyd.get_empty(len(df))
        for c in obsdata.dtype.names:
            assert c in df.columns, "Missing observation data field: {}".format(c)
            obsdata[c] = df[c]
        nhyd = len(df)
        hyd = flopy.modflow.ModflowHyd(self, nhyd=nhyd, hydnoh=-999, obsdata=obsdata)
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return hyd

    def setup_gag(self):

        print('setting up GAGE package...')
        t0 = time.time()
        # setup gage package output for all included lakes
        ngages = 0
        nlak_gages = 0
        starting_unit_number = self.cfg['gag']['starting_unit_number']
        if self.get_package('lak') is not None:
            nlak_gages = self.lak.nlakes
        if nlak_gages > 0:
            ngages += nlak_gages
            lak_gagelocs = list(np.arange(1, nlak_gages+1) * -1)
            lak_gagerch = [0] * nlak_gages # dummy list to maintain index position
            lak_outtype = [self.cfg['gag']['lak_outtype']] * nlak_gages
            # need minus sign to tell MF to read outtype
            lake_unit = list(-np.arange(starting_unit_number,
                                        starting_unit_number + nlak_gages, dtype=int))
            # TODO: make private attribute to facilitate keeping track of lake IDs
            lak_files = ['lak{}_{}.ggo'.format(i+1, hydroid)
                         for i, hydroid in enumerate(self.cfg['lak']['source_data']['lakes_shapefile']['include_ids'])]

        # need to add streams at some point
        nstream_gages = 0
        stream_gageseg = []
        stream_gagerch = []
        stream_unit = []
        stream_outtype = []
        stream_files = []
        if self.get_package('sfr') is not None:
            #observations_input = self.cfg['sfr'].get('source_data', {}).get('observations')
            #obs_info_files = self.cfg['gag'].get('observation_data')
            #if obs_info_files is not None:
            #    # get obs_info_files into dictionary format
            #    # filename: dict of column names mappings
            #    if isinstance(obs_info_files, str):
            #        obs_info_files = [obs_info_files]
            #    if isinstance(obs_info_files, list):
            #        obs_info_files = {f: self.cfg['gag']['default_columns']
            #                          for f in obs_info_files}
            #    elif isinstance(obs_info_files, dict):
            #        for k, v in obs_info_files.items():
            #            if v is None:
            #                obs_info_files[k] = self.cfg['gag']['default_columns']
#
            #    print('Reading observation files...')
            #    check_source_files(obs_info_files.keys())
            #    dfs = []
            #    for f, column_info in obs_info_files.items():
            #        print(f)
            #        df = read_observation_data(f,
            #                                   column_info,
            #                                   column_mappings=self.cfg['hyd'].get('column_mappings'))
            #        dfs.append(df) # cull to cols that are needed
            #    df = pd.concat(dfs, axis=0)
            df = self.sfrdata.observations
            nstream_gages = len(df)
            stream_files = ['{}.ggo'.format(site_no) for site_no in df.obsname]
            stream_gageseg = df.iseg.tolist()
            stream_gagerch = df.ireach.tolist()
            stream_unit = list(np.arange(starting_unit_number,
                                         starting_unit_number + nstream_gages, dtype=int))
            stream_outtype = [self.cfg['gag']['sfr_outtype']] * nstream_gages
        ngages += nstream_gages

        if ngages == 0:
            print('No gage package input.')
            return

        # create flopy gage package object
        gage_data = fm.ModflowGage.get_empty(ncells=ngages)
        gage_data['gageloc'] = lak_gagelocs + stream_gageseg
        gage_data['gagerch'] = lak_gagerch + stream_gagerch
        gage_data['unit'] = lake_unit + stream_unit
        gage_data['outtype'] = lak_outtype + stream_outtype
        if len(self.cfg['gag'].get('ggo_files', {})) == 0:
            self.cfg['gag']['ggo_files'] = lak_files + stream_files
        gag = fm.ModflowGage(self, numgage=len(gage_data),
                             gage_data=gage_data,
                             files=self.cfg['gag']['ggo_files'],
                             )
        print("finished in {:.2f}s\n".format(time.time() - t0))
        return gag

    def setup_chd(self):
        """
        Sets up the CHD package.

        Parameters
        ----------

        Notes
        -----

        """
        package = 'chd'
        print('\nSetting up {} package...'.format(package.upper()))
        t0 = time.time()
        package_config = self.cfg[package]

        # option to write stress_period_data to external files
        external_files = False  # not yet supported for MODFLOW-NWT
        external_filename_fmt = package_config.get('external_filename_fmt')

        # perimeter boundary
        if 'perimeter_boundary' in package_config:
            perimeter_cfg = package_config['perimeter_boundary']
            perimeter_cfg['boundary_type'] = 'head'
            if 'inset_parent_period_mapping' not in perimeter_cfg:
                perimeter_cfg['inset_parent_period_mapping'] = self.parent_stress_periods
            if 'parent_start_time' not in perimeter_cfg:
                perimeter_cfg['parent_start_date_time'] = self.parent.perioddata['start_datetime'][0]
            self.tmr = TmrNew(self.parent, self, **perimeter_cfg)
            perimeter_df = self.tmr.get_inset_boundary_values()

            # get the stress period data
            # this also sets up the external file paths
            spd = setup_flopy_stress_period_data(self, package, perimeter_df,
                                                 flopy_package_class=fm.ModflowChd,
                                                 variable_column='head',
                                                 external_files=external_files,
                                                 external_filename_fmt=external_filename_fmt)

        # placeholder for setting up user-specified CHD cells from CSV data
        # todo: support for non-perimeter chd cells
        df = pd.DataFrame()  # insert function here to get csv data into dataframe
        if len(df) == 0:
            print('No other CHD input specified')
            if 'perimeter_boundary' not in package_config:
                return

        kwargs = self.cfg[package].copy()
        if not external_files:
            kwargs['stress_period_data'] = spd

        kwargs = get_input_arguments(kwargs, fm.ModflowChd)
        chd = fm.ModflowChd(self, **kwargs)
        print("setup of chd took {:.2f}s\n".format(time.time() - t0))
        return chd

    def write_input(self):
        """Write the model input.
        """
        # prior to writing output
        # remove any BCs in inactive cells
        pckgs = ['CHD']
        for pckg in pckgs:
            package_instance = getattr(self, pckg.lower(), None)
            if package_instance is not None:
                remove_inactive_bcs(package_instance)

        # write the model with flopy
        # but skip the sfr package
        # by monkey-patching the write method
        SelPackList = [p for p in self.get_package_list() if p != 'SFR']
        super().write_input(SelPackList=SelPackList)

        # write the sfr package with SFRmaker
        # gage package was already set-up and then written by Flopy
        if 'SFR' in self.get_package_list():
            self.sfrdata.write_package(write_observations_input=False)

        # add version info to file headers
        files = [self.namefile]
        files += [p.file_name[0] for p in self.packagelist]
        for f in files:
            # either flopy or modflow
            # doesn't allow headers for some packages
            ext = Path(f).suffix
            if ext in {'.hyd', '.gag', '.gage'}:
                continue
            add_version_to_fileheader(f, model_info=self.header)

    @staticmethod
    def _parse_model_kwargs(cfg):
        return cfg

    @classmethod
    def load(cls, yamlfile, load_only=None, verbose=False, forgive=False, check=False):
        """Load a model from a config file and set of MODFLOW files.
        """
        cfg = load_cfg(yamlfile, verbose=verbose, default_file=cls.default_file) # '/mfnwt_defaults.yml')
        print('\nLoading {} model from data in {}\n'.format(cfg['model']['modelname'], yamlfile))
        t0 = time.time()

        m = cls(cfg=cfg, **cfg['model'])
        if 'grid' not in m.cfg.keys():
            m.setup_grid()
            #grid_file = cfg['setup_grid']['output_files']['grid_file']
            #if os.path.exists(grid_file):
            #    print('Loading model grid definition from {}'.format(grid_file))
            #    m.cfg['grid'] = load(grid_file)
            #else:
            #    m.setup_grid()

        m = flopy_mf2005_load(m, load_only=load_only, forgive=forgive, check=check)
        print('finished loading model in {:.2f}s'.format(time.time() - t0))
        return m


# -*-coding:UTF-8-*-
from scripts.logger.lemon_logger import Logger
from scripts.tools.mutator_selection_logic import Roulette, RandomMutant, MCMC, RandomMutator
import argparse
import sys
import ast
import os
import numpy as np
from itertools import combinations
import redis
import pickle
from scripts.tools import utils
import shutil
import re
import datetime
import configparser
import warnings
import math

np.random.seed(20200501)
warnings.filterwarnings("ignore")
os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ""

def partially_nan_or_inf(predictions, bk_num):
    """
    Check if there is NAN in the result
    """
    def get_nan_num(nds):
        _nan_num = 0
        for nd in nds:
            if np.isnan(nd).any() or np.isinf(nd).any():
                _nan_num += 1
        return _nan_num

    if len(predictions) == bk_num:
        for input_predict in zip(*predictions):
            nan_num = get_nan_num(input_predict)
            if 0 < nan_num < bk_num:
                return True
            else:
                continue
        return False
    else:
        raise Exception("wrong backend amounts")

def get_selector_by_startegy_name(mutator_s, mutant_s):

    mutant_strategy_dict = {'RANDOM':RandomMutant,"ROULETTE":Roulette}
    mutator_strategy_dict = {"RANDOM":RandomMutator,"MCMC":MCMC}
    return mutator_strategy_dict[mutator_s],mutant_strategy_dict[mutant_s]

def save_mutate_history(selector, invalid_history: dict, mutant_history: list):
    mutator_history_path = os.path.join(experiment_dir, "mutator_history.csv")
    mutant_history_path = os.path.join(experiment_dir, "mutant_history.txt")
    with open(mutator_history_path, "w+") as fw:
        fw.write("Name,Success,Invalid,Total\n")
        for op in invalid_history.keys():
            mtrs = selector.mutators[op]
            invalid_cnt = invalid_history[op]
            fw.write("{},{},{},{}\n".format(op, mtrs.delta_bigger_than_zero, invalid_cnt, mtrs.total))
    with open(mutant_history_path, "w+") as fw:
        for mutant in mutant_history:
            fw.write("{}\n".format(mutant))

def is_nan_or_inf(t):
    if math.isnan(t) or math.isinf(t):
        return True
    else:
        return False

def continue_checker(**run_stat):
    start_time = run_stat['start_time']
    time_limitation = run_stat['time_limit']
    cur_counters = run_stat['cur_counters']
    counters_limit = run_stat['counters_limit']
    s_mode = run_stat['stop_mode']

    # if timing
    if s_mode == 'TIMING':
        hours, minutes, seconds = utils.ToolUtils.get_HH_mm_ss(datetime.datetime.now() - start_time)
        total_minutes = hours * 60 + minutes
        mutate_logger.info(f"INFO: Mutation progress: {total_minutes}/{time_limitation} Minutes!")
        if total_minutes < time_limitation:
            return True
        else:
            return False
    # if counters
    elif s_mode == 'COUNTER':
        if cur_counters < counters_limit:
            mutate_logger.info("INFO: Mutation progress {}/{}".format(cur_counters + 1, counters_limit))
            return True
        else:
            return False
    else:
        raise Exception(f"Error! Stop Mode {s_mode} not Found!")

def _generate_and_predict(res_dict, filename, mutate_num, mutate_ops, test_size, exp, backends):
    """
    Generate models using mutate operators and store them
    """
    mutate_op_history = {k: 0 for k in mutate_ops}
    mutate_op_invalid_history = {k: 0 for k in mutate_ops}
    mutant_history = []

    # get mutator selection strategy
    origin_model_name = "{}_origin0.h5".format(exp)
    origin_save_path = os.path.join(mut_dir, origin_model_name)
    mutator_selector_func, mutant_selector_func = get_selector_by_startegy_name(mutator_strategy,mutant_strategy)
    mutator_selector, mutant_selector = mutator_selector_func(mutate_ops), mutant_selector_func([origin_model_name],
                                                                                capacity=mutate_num + 1)

    shutil.copy(src=filename, dst=origin_save_path)
    origin_model_status, res_dict, accumulative_inconsistency, _ = get_model_prediction(res_dict,
                                                                                origin_save_path,
                                                                                origin_model_name, exp,
                                                                                test_size, backends)
    if not origin_model_status:
        mutate_logger.error(f"Origin model {exp} crashed on some backends! LEMON would skip it")
        sys.exit(-1)

    last_used_mutator = None
    last_inconsistency = accumulative_inconsistency
    mutant_counter = 0

    start_time = datetime.datetime.now()
    order_inconsistency_dict = {}
    run_stat = {'start_time': start_time, 'time_limit': time_limit, 'cur_counters': mutant_counter,
                'counters_limit': mutate_num, 'stop_mode': stop_mode}

    while continue_checker(**run_stat):
        picked_seed = utils.ToolUtils.select_mutant(mutant_selector)
        selected_op = utils.ToolUtils.select_mutator(mutator_selector, last_used_mutator=last_used_mutator)
        mutate_op_history[selected_op] += 1
        last_used_mutator = selected_op
        mutator = mutator_selector.mutators[selected_op]
        mutant = mutant_selector.mutants[picked_seed]

        new_seed_name = "{}-{}{}.h5".format(picked_seed[:-3], selected_op, mutate_op_history[selected_op])
        # seed name would not be duplicate
        if new_seed_name not in mutant_selector.mutants.keys():
            new_seed_path = os.path.join(mut_dir, new_seed_name)
            picked_seed_path = os.path.join(mut_dir, picked_seed)
            mutate_st = datetime.datetime.now()
            mutate_status = os.system("{}/lemon/bin/python -m  scripts.mutation.model_mutation_generators --model {} "
                                      "--mutate_op {} --save_path {} --mutate_ratio {}".format(python_prefix,
                                                                                picked_seed_path,
                                                                                selected_op,
                                                                                new_seed_path,
                                                                                flags.mutate_ratio))
            mutate_et = datetime.datetime.now()
            mutate_dt = mutate_et - mutate_st
            h, m, s = utils.ToolUtils.get_HH_mm_ss(mutate_dt)
            mutate_logger.info("INFO:Mutate Time Used on {} : {}h, {}m, {}s".format(selected_op, h, m, s))
            # mutation status code is successful
            if mutate_status == 0:
                mutant.selected += 1
                mutator.total += 1
                # execute this model on all platforms
                predict_status, res_dict, accumulative_inconsistency, model_outputs = \
                    get_model_prediction(res_dict, new_seed_path, new_seed_name, exp, test_size, backends)

                if predict_status:
                    mutant_history.append(new_seed_name)

                    if utils.ModelUtils.is_valid_model(inputs_backends=model_outputs,backends_nums=len(backends)):
                        delta = accumulative_inconsistency - last_inconsistency

                        if mutator_strategy == 'MCMC':
                            mutator.delta_bigger_than_zero = mutator.delta_bigger_than_zero + 1 \
                                if delta > 0 else mutator.delta_bigger_than_zero

                        if mutant_strategy == 'RANDOM' or \
                                (mutant_strategy == 'ROULETTE' and delta > 0):
                            # when size >= capacity:
                            # random_mutant & Roulette would drop one and add new one
                            if mutant_selector.is_full():
                                mutant_selector.pop_one_mutant()
                            mutant_selector.add_mutant(new_seed_name)
                            last_inconsistency = accumulative_inconsistency

                        mutate_logger.info("SUCCESS:{} pass testing!".format(new_seed_name))
                        mutant_counter += 1
                    else:
                        mutate_op_invalid_history[selected_op] += 1
                        mutate_logger.error("Invalid model Found!")
                else:
                    mutate_logger.error("Crashed or NaN model Found!")
            else:
                mutate_logger.error("Exception raised when mutate {} with {}".format(picked_seed, selected_op))

            mutate_logger.info("Mutated op used history:")
            mutate_logger.info(mutate_op_history)

            mutate_logger.info("Invalid mutant generated history:")
            mutate_logger.info(mutate_op_invalid_history)

        run_stat['cur_counters'] = mutant_counter

    save_mutate_history(mutator_selector, mutate_op_invalid_history, mutant_history)

    return res_dict

def generate_metrics_result(res_dict, predict_output, model_idntfr):
    mutate_logger.info("Generating Metrics Result")
    accumulative_incons = 0
    backends_pairs_num = 0
    # Compare results pair by pair
    for pair in combinations(predict_output.items(), 2):
        backends_pairs_num += 1
        backend1, backend2 = pair
        bk_name1, prediction1 = backend1
        bk_name2, prediction2 = backend2
        bk_pair = "{}_{}".format(bk_name1, bk_name2)
        for metrics_name, metrics_result_dict in res_dict.items():
            metrics_func = utils.MetricsUtils.get_metrics_by_name(metrics_name)
            # metrics_results in list type
            metrics_results = metrics_func(prediction1, prediction2, y_test[:flags.test_size])

            # ACC -> float: The sum of all inputs under all backends
            accumulative_incons += sum(metrics_results)

            for input_idx, delta in enumerate(metrics_results):
                delta_key = "{}_{}_{}_input{}".format(model_idntfr, bk_name1, bk_name2, input_idx)
                metrics_result_dict[delta_key] = delta

    mutate_logger.info(f"Accumulative Inconsistency: {accumulative_incons}")
    return res_dict, accumulative_incons

def get_model_prediction(res_dict, model_path, model_name, exp, test_size, backends):
    """
    Get model prediction on different backends and calculate distance by metrics
    """
    predict_output = {b: [] for b in backends}
    model_idntfr = model_name[:-3]
    all_backends_predict_status = True
    for bk in backends:
        python_bin = f"{python_prefix}/{bk}/bin/python"
        predict_st = datetime.datetime.now()
        pre_status_bk = os.system(f"{python_bin} -u -m run.patch_prediction_extractor --backend {bk} "
                                  f"--exp {exp} --test_size {test_size} --model {model_path} "
                                  f"--redis_db {lemon_cfg['redis'].getint('redis_db')} --config_name {flags.config_name}")
        predict_et = datetime.datetime.now()
        predict_td = predict_et - predict_st
        h, m, s = utils.ToolUtils.get_HH_mm_ss(predict_td)
        mutate_logger.info("Prediction Time Used on {} : {}h, {}m, {}s".format(bk, h, m, s))

        # If no exception is thrown,save prediction result
        if pre_status_bk == 0:
            data = pickle.loads(redis_conn.hget("prediction_{}".format(model_name), bk))
            predict_output[bk] = data
        # record the crashed backend
        else:
            all_backends_predict_status = False
            mutate_logger.error("{} crash on backend {} when predicting ".format(model_name, bk))

    status = False
    accumulative_incons = None

    # run ok on all platforms
    if all_backends_predict_status:
        predictions = list(predict_output.values())
        res_dict, accumulative_incons = generate_metrics_result(res_dict=res_dict, predict_output=predict_output, model_idntfr=model_idntfr)

        # If all backends are working fine, check if there exists NAN or INF in the result
        # `accumulative_incons` is nan or inf --> NaN or INF in results
        if is_nan_or_inf(accumulative_incons):
            # has NaN on partial backends
            if partially_nan_or_inf(predictions, len(backends)):
                nan_model_path = os.path.join(nan_dir, f"{model_idntfr}_NaN_bug.h5")
                mutate_logger.error("Error: Found one NaN bug. move NAN model")

            # has NaN on all backends --> not a NaN bug
            else:
                nan_model_path = os.path.join(nan_dir, f"{model_idntfr}_NaN_on_all_backends.h5")
                mutate_logger.error("Error: Found one NaN Model on all libraries. move NAN model")
            shutil.move(model_path, nan_model_path)

        else:  # No NaN or INF on any backend
            mutate_logger.info("Saving prediction")
            with open("{}/prediction_{}.pkl".format(inner_output_dir, model_idntfr), "wb+") as f:
                pickle.dump(predict_output, file=f)
            status = True

    # save crashed model
    else:
        mutate_logger.error("Error: move crash model")
        crash_model_path = os.path.join(crash_dir, model_name)
        shutil.move(model_path, crash_model_path)

    return status, res_dict, accumulative_incons, predict_output

if __name__ == "__main__":

    starttime = datetime.datetime.now()
    """
    Parser of command args.
    It could make mutate_lemon.py run independently without relying on mutation_executor.py
    """
    parse = argparse.ArgumentParser()
    parse.add_argument("--is_mutate", type=ast.literal_eval, default=False,
                       help="parameter to determine mutation option")
    parse.add_argument("--mutate_op", type=str, nargs='+',
                       choices=['WS', 'GF', 'NEB', 'NAI', 'NS', 'ARem', 'ARep', 'LA', 'LC', 'LR', 'LS', 'MLA']
                       , help="parameter to determine mutation option")
    parse.add_argument("--model", type=str, help="relative path of model file(from root dir)")
    parse.add_argument("--output_dir", type=str, help="relative path of output dir(from root dir)")
    parse.add_argument("--backends", type=str, nargs='+', help="list of backends")
    parse.add_argument("--mutate_num", type=int, help="number of variant models generated by each mutation operator")
    parse.add_argument("--mutate_ratio", type=float, help="ratio of mutation")
    parse.add_argument("--exp", type=str, help="experiments identifiers")
    parse.add_argument("--test_size", type=int, help="amount of testing image")
    parse.add_argument("--config_name", type=str, help="config name")
    flags, unparsed = parse.parse_known_args(sys.argv[1:])

    warnings.filterwarnings("ignore")
    lemon_cfg = configparser.ConfigParser()
    lemon_cfg.read(f"./config/{flags.config_name}")
    time_limit = lemon_cfg['parameters'].getint("time_limit")
    mutator_strategy = lemon_cfg['parameters'].get("mutator_strategy").upper()
    mutant_strategy = lemon_cfg['parameters'].get("mutant_strategy").upper()
    stop_mode = lemon_cfg['parameters'].get("stop_mode").upper()
    alpha = lemon_cfg['parameters'].getfloat("alpha")

    mutate_logger = Logger()
    pool = redis.ConnectionPool(host=lemon_cfg['redis']['host'], port=lemon_cfg['redis']['port'],
                                db=lemon_cfg['redis'].getint('redis_db'))
    redis_conn = redis.Redis(connection_pool=pool)

    for k in redis_conn.keys():
        if flags.exp in k.decode("utf-8"):
            redis_conn.delete(k)

    # exp : like lenet5-mnist
    experiment_dir = os.path.join(flags.output_dir, flags.exp)
    mut_dir = os.path.join(experiment_dir, "mut_model")
    crash_dir = os.path.join(experiment_dir, "crash")
    nan_dir = os.path.join(experiment_dir, "nan")
    inner_output_dir = os.path.join(experiment_dir, "inner_output")
    metrics_result_dir = os.path.join(experiment_dir, "metrics_result")

    x, y = utils.DataUtils.get_data_by_exp(flags.exp)
    x_test, y_test = x[:flags.test_size], y[:flags.test_size]
    pool_size = lemon_cfg['parameters'].getint('pool_size')
    python_prefix = lemon_cfg['parameters']['python_prefix'].rstrip("/")

    try:
        metrics_list = lemon_cfg['parameters']['metrics'].split(" ")
        lemon_results = {k: dict() for k in metrics_list}
        lemon_results = _generate_and_predict(lemon_results, flags.model, flags.mutate_num, flags.mutate_op,
                                              flags.test_size, flags.exp, flags.backends)
        with open("{}/{}_lemon_results.pkl".format(experiment_dir, flags.exp), "wb+") as f:
            pickle.dump(lemon_results, file=f)
        utils.MetricsUtils.generate_result_by_metrics(metrics_list, lemon_results, metrics_result_dir, flags.exp)

    except Exception as e:
        mutate_logger.exception(sys.exc_info())

    from keras import backend as K
    K.clear_session()

    endtime = datetime.datetime.now()
    time_delta = endtime - starttime
    h, m, s = utils.ToolUtils.get_HH_mm_ss(time_delta)
    mutate_logger.info("Mutation process is done: Time used: {} hour,{} min,{} sec".format(h, m, s))


import numpy as np
import torch
import torch.nn as nn

from core.distributions import Bernoulli, Categorical, DiagGaussian
from core.utils import init
from core.agents.models import CNNBase, MLPBase

class Policy(nn.Module):
    def __init__(self, obs_shape, action_space, base=None, base_kwargs=None):
        super(Policy, self).__init__()
        # create base network for acting
        if base_kwargs is None:
            base_kwargs = {}
        if base is None:
            if len(obs_shape) == 3:
                base = CNNBase
            elif len(obs_shape) == 1:
                base = MLPBase
            else:
                raise NotImplementedError(
                    "Expect the observation's shape to be either 3 or 1 but got {}".format(len(obs_shape)))

        self.base = base(obs_shape[0], **base_kwargs)

        # action sampling mechanism
        if action_space.__class__.__name__ == "Discrete":
            num_outputs = action_space.n
            self.dist = Categorical(self.base.output_size, num_outputs)
        elif action_space.__class__.__name__ == "Box":
            num_outputs = action_space.shape[0]
            self.dist = DiagGaussian(self.base.output_size, num_outputs)
        elif action_space.__class__.__name__ == "MultiBinary":
            num_outputs = action_space.shape[0]
            self.dist = Bernoulli(self.base.output_size, num_outputs)
        else:
            raise NotImplementedError

        self.obs_shape = obs_shape
        self.action_space = action_space

    @property
    def is_recurrent(self):
        return self.base.is_recurrent

    @property
    def recurrent_hidden_state_size(self):
        """Size of rnn_hx."""
        return self.base.recurrent_hidden_state_size

    def forward(self, inputs, rnn_hxs, masks):
        raise NotImplementedError

    def act(self, inputs, rnn_hxs, masks, deterministic=False):
        value, actor_features, rnn_hxs = self.base(inputs, rnn_hxs, masks)
        dist = self.dist(actor_features)

        # if deterministic greedily choose the most optimal solution otherwise sampling with probability proportional to cummulate reward
        if deterministic:
            action = dist.mode()
        else:
            action = dist.sample()

        action_log_probs = dist.log_probs(action)
        dist_entropy = dist.entropy().mean()

        return value, action, action_log_probs, rnn_hxs

    def get_value(self, inputs, rnn_hxs, masks):
        value, _, _ = self.base(inputs, rnn_hxs, masks)
        return value

    def evaluate_actions(self, inputs, rnn_hxs, masks, action):
        value, actor_features, rnn_hxs = self.base(inputs, rnn_hxs, masks)
        dist = self.dist(actor_features)

        action_log_probs = dist.log_probs(action)
        dist_entropy = dist.entropy().mean()

        return value, action_log_probs, dist_entropy, rnn_hxs


# Copyright (c) 2021, Technische Universität Kaiserslautern (TUK) & National University of Sciences and Technology (NUST).
# All rights reserved.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

"""
    UNet model definition in here
"""

from __future__ import print_function
from __future__ import division
import torch
import torch.nn as nn
from torch.optim import *
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_
import torch.utils.model_zoo as model_zoo
from dataset import get_dataloaders_generated_data
import os
import numpy as np
import pickle as pkl
import PIL.Image as Image
import itertools
import matplotlib as mpl
# mpl.use('Agg')
import matplotlib.pyplot as plt
plt.switch_backend('agg')
from torchsummary import summary
from torchvision import models

# for getting pretrained layers from a vgg
matching_layers = [3, 6, 8, 11, 13, 16, 18]

class UNet_down_block(nn.Module):
    """
        Encoder class
    """
    def __init__(self, input_channel, output_channel, conv_1=None, conv_2=None):
        super(UNet_down_block, self).__init__()
        if conv_1:
            print('LOG: Using pretrained convolutional layer', conv_1)
        if conv_2:
            print('LOG: Using pretrained convolutional layer', conv_2)
        self.input_channels = input_channel
        self.output_channels = output_channel
        self.conv1 = conv_1 if conv_1 else nn.Conv2d(input_channel, output_channel, kernel_size=3, padding=1)
        self.conv2 = conv_2 if conv_2 else nn.Conv2d(output_channel, output_channel, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(num_features=output_channel)
        self.bn2 = nn.BatchNorm2d(num_features=output_channel)
        self.activate = nn.ReLU()

    def forward(self, x):
        x = self.activate(self.bn1(self.conv1(x)))
        x = self.activate(self.bn2(self.conv2(x)))
        return x

class UNet_up_block(nn.Module):
    """
        Decoder class
    """
    def __init__(self, prev_channel, input_channel, output_channel):
        super(UNet_up_block, self).__init__()
        self.output_channels = output_channel
        self.tr_conv_1 = nn.ConvTranspose2d(input_channel, input_channel, kernel_size=2, stride=2)
        self.conv_1 = nn.Conv2d(prev_channel+input_channel, output_channel, kernel_size=3, stride=1, padding=1)
        self.conv_2 = nn.Conv2d(output_channel, output_channel, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(num_features=output_channel)
        self.bn2 = nn.BatchNorm2d(num_features=output_channel)
        self.activate = nn.ReLU()

    def forward(self, prev_feature_map, x):
        x = self.tr_conv_1(x)
        x = self.activate(x)
        x = torch.cat((x, prev_feature_map), dim=1)
        x = self.activate(self.bn1(self.conv_1(x)))
        x = self.activate(self.bn2(self.conv_2(x)))
        return x

class UNet(nn.Module):
    def __init__(self, topology, input_channels, num_classes):
        super(UNet, self).__init__()
        # these topologies are possible right now
        self.topologies = {
            "ENC_1_DEC_1": self.ENC_1_DEC_1,
            "ENC_2_DEC_2": self.ENC_2_DEC_2,
            "ENC_3_DEC_3": self.ENC_3_DEC_3,
            "ENC_4_DEC_4": self.ENC_4_DEC_4,
        }
        assert topology in self.topologies
        vgg_trained = models.vgg11(pretrained=True)
        pretrained_layers = list(vgg_trained.features)
        self.max_pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout2d(0.6)
        self.activate = nn.ReLU()
        self.encoder_1 = UNet_down_block(input_channels, 64)
        self.encoder_2 = UNet_down_block(64, 128, conv_1=pretrained_layers[3])
        self.encoder_3 = UNet_down_block(128, 256, conv_1=pretrained_layers[6], conv_2=pretrained_layers[8])
        self.encoder_4 = UNet_down_block(256, 512, conv_1=pretrained_layers[11], conv_2=pretrained_layers[13])
        self.mid_conv_64_64_a = nn.Conv2d(64, 64, 3, padding=1)
        self.mid_conv_64_64_b = nn.Conv2d(64, 64, 3, padding=1)
        self.mid_conv_128_128_a = nn.Conv2d(128, 128, 3, padding=1)
        self.mid_conv_128_128_b = nn.Conv2d(128, 128, 3, padding=1)
        self.mid_conv_256_256_a = nn.Conv2d(256, 256, 3, padding=1)
        self.mid_conv_256_256_b = nn.Conv2d(256, 256, 3, padding=1)
        self.mid_conv_512_1024 = nn.Conv2d(512, 1024, 3, padding=1)
        self.mid_conv_1024_1024 = nn.Conv2d(1024, 1024, 3, padding=1)
        self.decoder_4 = UNet_up_block(prev_channel=self.encoder_4.output_channels, input_channel=self.mid_conv_1024_1024.out_channels, output_channel=256)
        self.decoder_3 = UNet_up_block(prev_channel=self.encoder_3.output_channels, input_channel=self.decoder_4.output_channels, output_channel=128)
        self.decoder_2 = UNet_up_block(prev_channel=self.encoder_2.output_channels, input_channel=self.decoder_3.output_channels, output_channel=64)
        self.decoder_1 = UNet_up_block(prev_channel=self.encoder_1.output_channels, input_channel=self.decoder_2.output_channels, output_channel=64)
        self.binary_last_conv = nn.Conv2d(64, num_classes, kernel_size=1)
        self.softmax = nn.Softmax(dim=1)
        self.forward = self.topologies[topology]
        print('\n\n' + "#" * 100)
        print("(LOG): The following Model Topology will be Utilized: {}".format(self.forward.__name__))
        print("#" * 100 + '\n\n')
        pass

    def ENC_1_DEC_1(self, x_in):
        x1_cat = self.encoder_1(x_in)
        x1_cat_1 = self.dropout(x1_cat)
        x1 = self.max_pool(x1_cat_1)
        x_mid = self.mid_conv_64_64_a(x1)
        x_mid = self.activate(x_mid)
        x_mid = self.mid_conv_64_64_b(x_mid)
        x_mid = self.activate(x_mid)
        x_mid = self.dropout(x_mid)
        x = self.decoder_1(x1_cat, x_mid)
        x = self.binary_last_conv(x)
        # return the final vector and the corresponding softmax-ed prediction
        return x, self.softmax(x)

    def ENC_2_DEC_2(self, x_in):
        x1_cat = self.encoder_1(x_in)
        x1 = self.max_pool(x1_cat)
        x2_cat = self.encoder_2(x1)
        x2_cat_1 = self.dropout(x2_cat)
        x2 = self.max_pool(x2_cat_1)
        x_mid = self.mid_conv_128_128_a(x2)
        x_mid = self.activate(x_mid)
        x_mid = self.mid_conv_128_128_b(x_mid)
        x_mid = self.activate(x_mid)
        x_mid = self.dropout(x_mid)
        x = self.decoder_2(x2_cat, x_mid)
        x = self.decoder_1(x1_cat, x)
        x = self.binary_last_conv(x)
        # return the final vector and the corresponding softmax-ed prediction
        return x, self.softmax(x)

    def ENC_3_DEC_3(self, x_in):
        x1_cat = self.encoder_1(x_in)
        x1 = self.max_pool(x1_cat)
        x2_cat = self.encoder_2(x1)
        x2_cat_1 = self.dropout(x2_cat)
        x2 = self.max_pool(x2_cat_1)
        x3_cat = self.encoder_3(x2)
        x3 = self.max_pool(x3_cat)
        x_mid = self.mid_conv_256_256_a(x3)
        x_mid = self.activate(x_mid)
        x_mid = self.mid_conv_256_256_b(x_mid)
        x_mid = self.activate(x_mid)
        x_mid = self.dropout(x_mid)
        x = self.decoder_3(x3_cat, x_mid)
        x = self.decoder_2(x2_cat, x)
        x = self.decoder_1(x1_cat, x)
        x = self.binary_last_conv(x)
        # return the final vector and the corresponding softmax-ed prediction
        return x, self.softmax(x)

    def ENC_4_DEC_4(self, x_in):
        x1_cat = self.encoder_1(x_in)
        x1 = self.max_pool(x1_cat)
        x2_cat = self.encoder_2(x1)
        x2_cat_1 = self.dropout(x2_cat)
        x2 = self.max_pool(x2_cat_1)
        x3_cat = self.encoder_3(x2)
        x3 = self.max_pool(x3_cat)
        x4_cat = self.encoder_4(x3)
        x4_cat_1 = self.dropout(x4_cat)
        x4 = self.max_pool(x4_cat_1)
        x_mid = self.mid_conv_512_1024(x4)
        x_mid = self.activate(x_mid)
        x_mid = self.mid_conv_1024_1024(x_mid)
        x_mid = self.activate(x_mid)
        x_mid = self.dropout(x_mid)
        x = self.decoder_4(x4_cat, x_mid)
        x = self.decoder_3(x3_cat, x)
        x = self.decoder_2(x2_cat, x)
        x = self.decoder_1(x1_cat, x)
        x = self.binary_last_conv(x)
        # return the final vector and the corresponding softmax-ed prediction
        return x, self.softmax(x)
    pass

@torch.no_grad()
def check_model(topology, input_channels, num_classes, input_shape):
    model = UNet(topology=topology, input_channels=input_channels, num_classes=num_classes)
    model.eval()
    in_tensor = torch.Tensor(*input_shape)
    with torch.no_grad():
        out_tensor, softmaxed = model(in_tensor)
        print(in_tensor.shape, out_tensor.shape)
    pass

@torch.no_grad()
def check_model_on_dataloader():
    model = UNet(input_channels=11, num_classes=16)
    model.eval()
    model.cuda(device=0)
    loaders = get_dataloaders_generated_data(generated_data_path='generated_dataset', save_data_path='pickled_generated_datalist.pkl', block_size=256,
                                             model_input_size=64, batch_size=128, num_workers=8)
    with torch.no_grad():
        train_dataloader, val_dataloader, test_dataloader = loaders
        for idx, data in enumerate(train_dataloader):
            examples, labels = data['input'], data['label']
            examples = examples.cuda(device=0)
            print('-> on batch {}/{}, {}'.format(idx + 1, len(train_dataloader), examples.size()))
            out_tensor, prediction = model(examples)
            print(examples.shape, labels.shape, out_tensor.shape,
                  prediction.shape, torch.argmax(prediction, dim=1)[0,:,:].shape)

    pass

def see_children_recursively(graph, layer=None):
    further = False
    children = list(graph.children())
    for child in children:
        further = True
        if layer:
            see_children_recursively(child, layer)
        else:
            see_children_recursively(child)
    if layer:
        if not further and isinstance(graph, layer):
            print(graph)
    else:
        if not further:
            print(graph)
    # matching_pairs.append((graph.in_channels, graph.out_channels))
    pass

if __name__ == '__main__':
    # check_model
    check_model(topology="ENC_1_DEC_1", input_channels=7, num_classes=2, input_shape=[4, 7, 64, 64])
    # model = UNet(topology='ENC_1_DEC_1', input_channels=3, num_classes=4)
    # # model.cuda(device='cuda')
    # model.eval()
    # with torch.no_grad():
    #     summary(model, input_size=(3, 128, 128))
    # see_children_recursively(graph=model, layer=nn.Conv2d)
    pass


# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import time
import logging
import os
from pycocotools.coco import COCO

import numpy as np
import torch
import torchvision
import random
import math
import cv2
import glob as _glob
import csv

from core.evaluate import accuracy
from core.inference import get_final_preds
from utils.transforms import flip_back
from utils.vis import save_debug_images

import gc
import json

logger = logging.getLogger(__name__)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#시드고정
torch.manual_seed(0)
if device == 'cuda':
    torch.cuda.manual_seed_all(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(1)
random.seed(1)

csv.register_dialect(
    'mydialect',
    delimiter = ',',
    quotechar = '"',
    doublequote = True,
    skipinitialspace = True,
    lineterminator = '\r\n',
    quoting = csv.QUOTE_MINIMAL)

def writecsv(csvname,contents):
    f = open(csvname, 'a', newline='')
    wr = csv.writer(f)
    wr.writerow(contents)
    f.close()

def glob(dir, pats, recursive=False):  # faster than match, python3 only
    pats = pats if isinstance(pats, (list, tuple)) else [pats]
    matches = []
    for pat in pats:
        matches += _glob.glob(os.path.join(dir, pat), recursive=recursive)
    return matches

def train(config, train_loader, model, criterion, optimizer, epoch,
          output_dir, tb_log_dir, writer_dict):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    acc = AverageMeter()

    # switch to train mode
    model.train()
    model=model.to(device)

    end = time.time()
    for i, (input, target, target_weight, meta) in enumerate(train_loader):

        # measure data loading time
        data_time.update(time.time() - end)

        input=input.type(torch.FloatTensor)
        input=input.to(device)
        # compute output
        outputs = model(input)

        target = target.cuda(non_blocking=True)
        target_weight = target_weight.cuda(non_blocking=True)

        if isinstance(outputs, list):
            loss = criterion(outputs[0], target, target_weight)
            for output in outputs[1:]:
                loss += criterion(output, target, target_weight)
        else:
            output = outputs
            loss = criterion(output, target, target_weight)

        # loss = criterion(output, target, target_weight)

        # compute gradient and do update step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # measure accuracy and record loss
        losses.update(loss.item(), input.size(0))

        _, avg_acc, cnt, pred = accuracy(output.detach().cpu().numpy(),
                                         target.detach().cpu().numpy())
        acc.update(avg_acc, cnt)

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % config.PRINT_FREQ == 0:
            msg = 'Epoch: [{0}][{1}/{2}]\t' \
                  'Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t' \
                  'Speed {speed:.1f} samples/s\t' \
                  'Data {data_time.val:.3f}s ({data_time.avg:.3f}s)\t' \
                  'Loss {loss.val:.5f} ({loss.avg:.5f})\t' \
                  'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
                      epoch, i, len(train_loader), batch_time=batch_time,
                      speed=input.size(0)/batch_time.val,
                      data_time=data_time, loss=losses, acc=acc)
            logger.info(msg)

            writer = writer_dict['writer']
            global_steps = writer_dict['train_global_steps']
            writer.add_scalar('train_loss', losses.val, global_steps)
            writer.add_scalar('train_acc', acc.val, global_steps)
            writer_dict['train_global_steps'] = global_steps + 1

            prefix = '{}_{}'.format(os.path.join(output_dir, 'train'), i)
            save_debug_images(config, input, meta, target, pred*4, output,
                              prefix)
        random.seed(1)

def validate(config, val_loader, val_dataset, model, criterion, output_dir,
             tb_log_dir, writer_dict=None):
    batch_time = AverageMeter()
    losses = AverageMeter()
    acc = AverageMeter()

    # switch to evaluate mode
    model.eval()
    model=model.to(device)

    num_samples = len(val_dataset)
    all_preds = np.zeros(
        (num_samples, config.MODEL.NUM_JOINTS, 3),
        dtype=np.float32
    )
    all_boxes = np.zeros((num_samples, 6))
    image_path = []
    filenames = []
    imgnums = []
    idx = 0
    with torch.no_grad():
        end = time.time()
        for i, (input, target, target_weight, meta) in enumerate(val_loader):
            input=input.to(device)

            # compute output
            outputs = model(input)
            if isinstance(outputs, list):
                output = outputs[-1]
            else:
                output = outputs

            if config.TEST.FLIP_TEST:
                # this part is ugly, because pytorch has not supported negative index
                # input_flipped = model(input[:, :, :, ::-1])
                input_flipped = np.flip(input.cpu().numpy(), 3).copy()
                input_flipped = torch.from_numpy(input_flipped).cuda()
                outputs_flipped = model(input_flipped)

                if isinstance(outputs_flipped, list):
                    output_flipped = outputs_flipped[-1]
                else:
                    output_flipped = outputs_flipped

                output_flipped = flip_back(output_flipped.cpu().numpy(),
                                           val_dataset.flip_pairs)
                output_flipped = torch.from_numpy(output_flipped.copy()).cuda()

                # feature is not aligned, shift flipped heatmap for higher accuracy
                if config.TEST.SHIFT_HEATMAP:
                    output_flipped[:, :, :, 1:] = \
                        output_flipped.clone()[:, :, :, 0:-1]

                output = (output + output_flipped) * 0.5

            target = target.cuda(non_blocking=True)
            target_weight = target_weight.cuda(non_blocking=True)

            loss = criterion(output, target, target_weight)

            num_images = input.size(0)
            # measure accuracy and record loss
            losses.update(loss.item(), num_images)
            _, avg_acc, cnt, pred = accuracy(output.cpu().numpy(),
                                             target.cpu().numpy())

            acc.update(avg_acc, cnt)

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            c = meta['center'].numpy()
            s = meta['scale'].numpy()
            score = meta['score'].numpy()

            preds, maxvals = get_final_preds(
                config, output.clone().cpu().numpy(), c, s)

            all_preds[idx:idx + num_images, :, 0:2] = preds[:, :, 0:2]
            all_preds[idx:idx + num_images, :, 2:3] = maxvals
            # double check this all_boxes parts
            all_boxes[idx:idx + num_images, 0:2] = c[:, 0:2]
            all_boxes[idx:idx + num_images, 2:4] = s[:, 0:2]
            all_boxes[idx:idx + num_images, 4] = np.prod(s*200, 1)
            all_boxes[idx:idx + num_images, 5] = score
            image_path.extend(meta['image'])

            idx += num_images

            if i % config.PRINT_FREQ == 0:
                msg = 'Test: [{0}/{1}]\t' \
                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
                      'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
                          i, len(val_loader), batch_time=batch_time,
                          loss=losses, acc=acc)
                logger.info(msg)

                prefix = '{}_{}'.format(
                    os.path.join(output_dir, 'val'), i
                )
                save_debug_images(config, input, meta, target, pred*4, output,
                                  prefix)
        '''
        name_values, perf_indicator = val_dataset.evaluate(
            config, all_preds, output_dir, all_boxes, image_path,
            filenames, imgnums
        )
        '''
        model_name = config.MODEL.NAME
        '''
        if isinstance(name_values, list):
            for name_value in name_values:
                _print_name_value(name_value, model_name)
        else:
            _print_name_value(name_values, model_name)
        '''
        if writer_dict:
            writer = writer_dict['writer']
            global_steps = writer_dict['valid_global_steps']
            writer.add_scalar(
                'valid_loss',
                losses.avg,
                global_steps
            )
            writer.add_scalar(
                'valid_acc',
                acc.avg,
                global_steps
            )
            '''
            if isinstance(name_values, list):
                for name_value in name_values:
                    writer.add_scalars(
                        'valid',
                        dict(name_value),
                        global_steps
                    )
            else:
                writer.add_scalars(
                    'valid',
                    dict(name_values),
                    global_steps
                )
            '''
            writer_dict['valid_global_steps'] = global_steps + 1

    return losses.avg

def test(config, val_loader, val_dataset, model,  output_dir,TTA,TTA_VALT,start):
    batch_time = AverageMeter()

    # switch to evaluate mode
    model.eval()
    model=model.to(device)

    num_samples = len(val_dataset)
    all_preds = np.zeros(
        (num_samples, config.MODEL.NUM_JOINTS, 3),
        dtype=np.float32
    )
    all_boxes = np.zeros((num_samples, 6))
    image_path = []
    filenames = []
    imgnums = []
    idx = 0

    file_path='./datasets/sample_submission.json'
    json_data = {}
    with open(file_path, "r") as json_file:
        json_data = json.load(json_file)

    #### 좌표 얼마나 차이나는지 확인용 실제 test 데이터 inference시에는 주석처리  #1
    annos=glob('datasets/train/annotations','*')
    annos.sort()
    annos=annos[8000:]
    #### 좌표 얼마나 차이나는지 확인용 실제 test 데이터 inference시에는 주석처리  #1

    with torch.no_grad():

        for i, (input, target, target_weight, center,scale,score,affine,origin_img,filename) in enumerate(val_loader):
            input=input.type(torch.FloatTensor)

            if TTA==True:
                input=input[0].to(device)
            else:
                input=input.to(device)
            # compute output
            outputs = model(input)
            if isinstance(outputs, list):
                output = outputs[-1]
            else:
                output = outputs

            num_images = input.size(0)

            c = center.numpy()
            s = scale.numpy()
            score = score.numpy()

            # 예측값 웝본 사이즈에 맞춰서 다시 scaling 처리
            if TTA==True:
                pred,_=get_max_preds_only_preds(output.cpu().numpy())

            else:
                pred,_=get_max_preds_only_preds(output.cpu().numpy())

            ## 좌표 반환 및 실제 사이즈가 아닌 input 사이즈 대비 좌표 반환
            if TTA==True:
                if TTA_VALT=='AVG':
                    pred=np.expand_dims(np.mean(pred,axis=0),axis=0)
                    points = save_batch_image_only_preds(input[0:1,:,:,:],pred*4,'test'+str(i)+'.jpg',save=False)

                elif TTA_VALT=='MEDIAN':
                    points=[]
                    for pred_val in pred:
                        points.append(save_batch_image_only_preds(input[0:1,:,:,:],np.expand_dims(pred_val,axis=0)*4,'test'+str(i)+'.jpg',save=False))
                    npdata = np.array(points)
                    npdata = np.median(npdata,axis=0)
                    points = npdata

            else:
                points = save_batch_image_only_preds(input,pred*4,'test'+str(i)+'.jpg',save=False)

            #### 좌표 얼마나 차이나는지 확인용 실제 test 데이터 inference시에는 주석처리  #2
            annot=COCO(annos[i])

            cow_annot=np.array(annot.dataset['label_info']['annotations'][0]['keypoints'])
            cow_annot=np.reshape(cow_annot,newshape=(17,3))[:17,:2]

            img=cv2.imread('datasets/train/images/'+annot.dataset['label_info']['image']['file_name'], cv2.IMREAD_COLOR)
            #### 좌표 얼마나 차이나는지 확인용 실제 test 데이터 inference시에는 주석처리  #2

            ## 실제 training 시에는 padding을 포함한 affine translation 으로 이미지를 translation and resize를 진행하기 떄문에 마지막 output에서
            ## 이를 보정하여 좌표를 계산해야함

            ## 좌표계산
            affine=affine.numpy()
            points=np.array(points)

            if TTA==True:
                points[:,0]=(points[:,0]-int(affine[0][0][2])) * (origin_img[0].shape[2]/(input[0].shape[2]-int(affine[0][0][2])*2))
                points[:,1]=(points[:,1]-int(affine[0][1][2])) * (origin_img[0].shape[1]/(input[0].shape[1]-int(affine[0][1][2])*2))
            else:
                points[:,0]=(points[:,0]-int(affine[0][0][2])) * (origin_img.shape[2]/(input.shape[3]-int(affine[0][0][2])*2))
                points[:,1]=(points[:,1]-int(affine[0][1][2])) * (origin_img.shape[1]/(input.shape[2]-int(affine[0][1][2])*2))

            #### 좌표 얼마나 차이나는지 확인용 실제 test 데이터 inference시에는 주석처리  #3
            test_imgssss=input[0].permute(1,2,0).cpu().numpy()[int(affine[0][1][2]):,int(affine[0][0][2]):,:]
            #save_test(img,points,'./test_folds/CA_WU_CP_N_C_B/'+str(i)+'.jpg')
            diff=np.abs(np.sum(cow_annot-points)) #단순 좌표값 차이의 합계 계산
            jointbad=np.zeros(shape=17)
            #야메 PCK
            #기준 소의 전체 영역의 min max x,y  좌표를 가져와서 해당 좌표의 상하좌우 10% point 만큼 안에 들어오면 ok 안되면 fail로 xx/17로 score 계산
            #n% point 만큼 조금씩 조정하면서 해보기
            pck, out, tor=cal_pck(cow_annot,points,0.35)

            ls=json_data['annotations'][i]
            for idx in range(17):
                ls['joint_self'][idx]=points[idx].tolist()

            del output,input,points,pred
            gc.collect()
            #if pck<1:
            #    save_bad(img,points,cow_annot,'./test_folds/bad_pck_20210805_MIDVAL/'+filename[0].split('/')[-1])
            #    for idx in out:
            #        jointbad[idx]+=1

            writecsv('valids_flip_with_pck_joint_W_ORIGIN.csv',[filename[0],diff,pck,tor])
            #writecsv('joint_bad_W_MIDVAL.csv',jointbad.tolist())
            #### 좌표 얼마나 차이나는지 확인용 실제 test 데이터 inference시에는 주석처리  #3
    end = time.time()
    print(end-start)
    json_data['latency']=end-start
    output_path='output_test.json'
    with open(output_path, 'w') as outfile:
        json.dump(json_data, outfile)
    print('end')

def cal_pck(anno,point,ratio):
    # baseline 기준 torso는 3번과 10번의 y좌표의 차이를 기준으로 하기 떄문에 똑같이 처리함
    tor= np.linalg.norm(np.abs((anno[3][:]-anno[10][:])))*ratio

    diff=np.abs(point-anno)
    out=np.concatenate([np.where(diff[:,0]>tor)[0],np.where(diff[:,1]>tor)[0]])
    out=np.unique(out)
    pck=1
    if out.shape[0]>0:
        pck=(17-out.shape[0])/17.0

    return pck, out, tor

def save_bad(img,points,annos,file_name):
    for point in points:
        print((int(point[0]), int(point[1])))
        cv2.circle(img, (int(point[0]), int(point[1])), 2, [0, 0, 255], 2)

    for point in annos:
        print((int(point[0]), int(point[1])))
        cv2.circle(img, (int(point[0]), int(point[1])), 2, [0, 255, 0], 2)
    cv2.imwrite(file_name, img)

#단일 좌표들만 찍을때
def save_test(img,points,file_name):
    for point in points:
        cv2.circle(img, (int(point[0]), int(point[1])), 2, [255, 0, 0], 2)
    cv2.imwrite(file_name, img)

def save_batch_image_only_preds(batch_image, batch_joints,
                                 file_name, nrow=8, padding=2,save=True):
    '''
    batch_image: [batch_size, channel, height, width]
    batch_joints: [batch_size, num_joints, 3],
    batch_joints_vis: [batch_size, num_joints, 1],
    }
    '''
    grid = torchvision.utils.make_grid(batch_image, nrow, padding, True)
    ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
    ndarr = ndarr.copy()

    nmaps = batch_image.size(0)
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height = int(batch_image.size(2) + padding)
    width = int(batch_image.size(3) + padding)
    points =[]
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            joints = batch_joints[k]

            for joint in joints:
                joint[0] = x * width + padding + joint[0]
                joint[1] = y * height + padding + joint[1]

                cv2.circle(ndarr, (int(joint[0]), int(joint[1])), 2, [255, 0, 0], 2)
                points.append([int(joint[0]),int(joint[1])])
            k = k + 1
    if save==True:
        cv2.imwrite(file_name, ndarr)
    return points

def get_max_preds_only_preds(batch_heatmaps):
    '''
    get predictions from score maps
    heatmaps: numpy.ndarray([batch_size, num_joints, height, width])
    '''
    assert isinstance(batch_heatmaps, np.ndarray), \
        'batch_heatmaps should be numpy.ndarray'
    assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim'

    batch_size = batch_heatmaps.shape[0]
    num_joints = batch_heatmaps.shape[1]
    width = batch_heatmaps.shape[3]
    heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1))
    idx = np.argmax(heatmaps_reshaped, 2)
    maxvals = np.amax(heatmaps_reshaped, 2)

    maxvals = maxvals.reshape((batch_size, num_joints, 1))
    idx = idx.reshape((batch_size, num_joints, 1))

    preds = np.tile(idx, (1, 1, 2)).astype(np.float32)

    preds[:, :, 0] = (preds[:, :, 0]) % width
    preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)

    pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))
    pred_mask = pred_mask.astype(np.float32)

    preds *= pred_mask
    return preds, maxvals

# markdown format output
def _print_name_value(name_value, full_arch_name):
    names = name_value.keys()
    values = name_value.values()
    num_values = len(name_value)
    logger.info(
        '| Arch ' +
        ' '.join(['| {}'.format(name) for name in names]) +
        ' |'
    )
    logger.info('|---' * (num_values+1) + '|')

    if len(full_arch_name) > 15:
        full_arch_name = full_arch_name[:8] + '...'
    logger.info(
        '| ' + full_arch_name + ' ' +
        ' '.join(['| {:.3f}'.format(value) for value in values]) +
         ' |'
    )

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count if self.count != 0 else 0


import tensorflow as tf
import data_loader_recsys
import generator_recsys
import utils
import shutil
import time
import math
import eval
import numpy as np
import argparse

#check whether the files exists or not,   "Data/Models/generation_model/model_nextitnet.ckpt"
#  if yes run this file directly, if not run nextitrec.py first, which is the training file.

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--top_k', type=int, default=5,
                        help='Sample from top k predictions')
    parser.add_argument('--beta1', type=float, default=0.9,
                        help='hyperpara-Adam')

    parser.add_argument('--datapath', type=str, default='Data/Session/user-filter-20000items-session5.csv',
                        help='data path')
    parser.add_argument('--eval_iter', type=int, default=10,
                        help='Sample generator output evry x steps')
    parser.add_argument('--save_para_every', type=int, default=10,
                        help='save model parameters every')
    parser.add_argument('--tt_percentage', type=float, default=0.2,
                        help='default=0.2 means 80% training 20% testing')
    parser.add_argument('--is_generatesubsession', type=bool, default=False,
                        help='whether generating a subsessions, e.g., 12345-->01234,00123,00012  It may be useful for very some very long sequences')
    args = parser.parse_args()

    dl = data_loader_recsys.Data_Loader({'model_type': 'generator', 'dir_name': args.datapath})
    all_samples = dl.item
    items = dl.item_dict
    print "len(items)",len(items)

    # Randomly shuffle data
    np.random.seed(10)
    shuffle_indices = np.random.permutation(np.arange(len(all_samples)))
    all_samples = all_samples[shuffle_indices]

    # Split train/test set
    dev_sample_index = -1 * int(args.tt_percentage * float(len(all_samples)))
    train_set, valid_set = all_samples[:dev_sample_index], all_samples[dev_sample_index:]

    model_para = {
        #all parameters shuold be consist with those in nextitred.py!!!!
        'item_size': len(items),
        'dilated_channels': 100,
        'dilations': [1, 2,],
        'kernel_size': 3,
        'learning_rate':0.001,
        'batch_size':32,
        'iterations':2,#useless, can be removed
        'is_negsample':False #False denotes no negative sampling
    }

    itemrec = generator_recsys.NextItNet_Decoder(model_para)
    itemrec.train_graph(model_para['is_negsample'])
    itemrec.predict_graph(model_para['is_negsample'],reuse=True)

    sess= tf.Session()
    init=tf.global_variables_initializer()
    sess.run(init)

    saver = tf.train.Saver()
    saver.restore(sess,"Data/Models/generation_model/model_nextitnet.ckpt")

    batch_no_test = 0
    batch_size_test = model_para['batch_size']
    curr_preds_5 = []
    rec_preds_5 = []  # 1
    ndcg_preds_5 = []  # 1
    curr_preds_20 = []
    rec_preds_20 = []  # 1
    ndcg_preds_20 = []  # 1
    while (batch_no_test + 1) * batch_size_test < valid_set.shape[0]:
        item_batch = valid_set[batch_no_test * batch_size_test: (batch_no_test + 1) * batch_size_test, :]
        [probs] = sess.run(
            [itemrec.g_probs],
            feed_dict={
                itemrec.input_predict: item_batch
            })
        for bi in range(probs.shape[0]):
            pred_items_5 = utils.sample_top_k(probs[bi][-1], top_k=args.top_k)  # top_k=5
            pred_items_20 = utils.sample_top_k(probs[bi][-1], top_k=args.top_k + 15)

            true_item = item_batch[bi][-1]
            predictmap_5 = {ch: i for i, ch in enumerate(pred_items_5)}
            pred_items_20 = {ch: i for i, ch in enumerate(pred_items_20)}

            rank_5 = predictmap_5.get(true_item)
            rank_20 = pred_items_20.get(true_item)
            if rank_5 == None:
                curr_preds_5.append(0.0)
                rec_preds_5.append(0.0)  # 2
                ndcg_preds_5.append(0.0)  # 2
            else:
                MRR_5 = 1.0 / (rank_5 + 1)
                Rec_5 = 1.0  # 3
                ndcg_5 = 1.0 / math.log(rank_5 + 2, 2)  # 3
                curr_preds_5.append(MRR_5)
                rec_preds_5.append(Rec_5)  # 4
                ndcg_preds_5.append(ndcg_5)  # 4
            if rank_20 == None:
                curr_preds_20.append(0.0)
                rec_preds_20.append(0.0)  # 2
                ndcg_preds_20.append(0.0)  # 2
            else:
                MRR_20 = 1.0 / (rank_20 + 1)
                Rec_20 = 1.0  # 3
                ndcg_20 = 1.0 / math.log(rank_20 + 2, 2)  # 3
                curr_preds_20.append(MRR_20)
                rec_preds_20.append(Rec_20)  # 4
                ndcg_preds_20.append(ndcg_20)  # 4

        batch_no_test += 1
        print "BATCH_NO: {}".format(batch_no_test)
        print "Accuracy mrr_5:", sum(curr_preds_5) / float(len(curr_preds_5))  # 5
        print "Accuracy mrr_20:", sum(curr_preds_20) / float(len(curr_preds_20))  # 5
        print "Accuracy hit_5:", sum(rec_preds_5) / float(len(rec_preds_5))  # 5
        print "Accuracy hit_20:", sum(rec_preds_20) / float(len(rec_preds_20))  # 5
        print "Accuracy ndcg_5:", sum(ndcg_preds_5) / float(len(ndcg_preds_5))  # 5
        print "Accuracy ndcg_20:", sum(ndcg_preds_20) / float(len(ndcg_preds_20))  #
        # print "curr_preds",curr_preds

if __name__ == '__main__':
    main()