from __future__ import division
import re
import os
import sys
import math
import random
import numpy as np
import pandas as pd
from PyDAIR.seq.IgSeq import *
from PyDAIR.io.PyDAIRIO import *

class PyDAIRDiversity:
    """PyDAIRDiversity class.

    A class is used for saving the results of diversity analysis.
    """

    def __init__(self):
        """PyDAIRDiversity class initialize method.

        Set up `None` object in default.
        """

        self.rarefaction        = {'vdj': None, 'cdr3': None}
        self.samplingresampling = {'vdj': None, 'cdr3': None}

class PyDAIRStatsRecord:
    """PyDAIRStatsRecord class.

    One PYDAIR file should has one **BLGIHStatsRecord** class object. If there
    are more than one PYDAIR files, use a number of **PyDAIRStatsRecord** class
    objects to save these data, and save all **PyDAIRStatsRecord** class objects
    into **PyDAIRStatsRecords** class.
    This class requires the list of V gene names, D gene names, J gene names,
    CDR3 nucleotide and protein seuqence, and the stop codon tags.
    """

    def __init__(self, name = None, v = None, d = None, j = None, orf = None,
                 cdr3_nucl_seq = None, cdr3_prot_seq = None,
                 v_del = None, j_del = None, vj_ins = None,
                 discard_ambiguous_D = False, productive_only = False):
        """PyDAIRStatsRecord class initialize method.

        Args:
            name (str): Sample name.
            v (list): Assigned V gene names.
            d (list): Assigned D gene names.
            j (list): Assigned J gene names.
            orf (list): A list of ORF.
            cdr3_nucl_seq (list): A list of CDR3 nucleotide sequences.
            cdr3_prot_seq (list): A list of CDR3 amino acid sequences.
            v_del (list): Deleted nucleotides of 3'-end V gene.
            j_del (list): Deleted nucleotides of 5'-end J gene.
            vj_ins (list): Inserted nucleotides.
            discard_ambiguous_D (bool): Default is `False`. If `True`, the sequences
                                        with ambiguous D segment will be discarded before
                                        summarization.
            productive_only (bool): Default is `False`. If `True`, the sequences contained
                                    more than one stop codon will be discarded before
                                    summarization.

        """

        # set default data
        self.name = name
        self.vdj  = pd.DataFrame({'v': v, 'd': d, 'j': j}, columns = ['v', 'd', 'j']).fillna(value = np.nan)
        self.cdr3 = pd.DataFrame({'nucl_seq': cdr3_nucl_seq, 'prot_seq': cdr3_prot_seq,
                                  'nucl_len': pd.Series(cdr3_nucl_seq).str.len(),
                                  'prot_len': pd.Series(cdr3_prot_seq).str.len()},
                                  columns = ['nucl_seq', 'prot_seq', 'nucl_len', 'prot_len']).fillna(value = np.nan)
        self.indels = pd.DataFrame({'v_del': v_del, 'v_del_len': pd.Series(v_del).str.len(),
                                    'j_del': j_del, 'j_del_len': pd.Series(j_del).str.len(),
                                    'vj_ins': vj_ins, 'vj_ins_len': pd.Series(vj_ins).str.len()},
                                    columns = ['v_del', 'v_del_len', 'j_del', 'j_del_len',
                                               'vj_ins', 'vj_ins_len']).fillna(value = np.nan)

        # set filters
        filter_ambigoD = None
        if discard_ambiguous_D:
            filter_ambigoD = self.vdj.d.notnull()
        else:
            filter_ambigoD = pd.Series([True] * self.vdj.shape[0])

        filter_stopcodon = None
        if productive_only:
            filter_stopcodon = pd.Series(orf).notnull()
        else:
            filter_stopcodon = pd.Series([True] * self.vdj.shape[0])

        filters = pd.Series(filter_ambigoD & filter_stopcodon)

        # filter data
        self.vdj  = self.vdj[filters]
        self.cdr3 = self.cdr3[filters]
        self.indels = self.indels[filters]

        # diversity study
        self.div = PyDAIRDiversity()

    def __len__(self):
        return self.vdj.shape[0]

    def len(self):
        """Retrive the number of sequences analyzed.

        """

        return self.__len__()

    def get_summary(self, data_type, prob = False, func = 'mean'):
        """Returns summary statistics.

        Args:
            data_type (str): A string to specify data type. ``v``,
                             ``d``, ``j``, ``vdj``, ``cdr3_nucl_len``,
                             ``cdr3_prot_len``, ``v_del_len``, ``j_del_len``,
                             and ``vj_ins_len`` are supported.
            prob (bool): If ``True``, calculate the probability.
            func (str): A string to specify data calculations when data_type is vdj_rarefation.
        """

        s = None

        if data_type in ['v', 'd', 'j', 'cdr3_prot_len', 'cdr3_nucl_len', 'v_del_len', 'j_del_len', 'vj_ins_len']:
            s = self.__get_freq(data_type, prob)
        elif data_type == 'vdj':
            s = self.__get_freq_vdj(data_type, prob)
        elif data_type == 'vdj_rarefaction':
            s = self.__get_est_vdj_rarefaction(data_type, func)
        else:
            raise ValueError(data_type + ' is not supported by get_summary method.')

        return s

    def __get_freq(self, data_type, prob = False):
        freq = None
        if data_type == 'v_del_len':
            freq = self.indels.v_del_len.value_counts(dropna = False)
            freq.index = ['ambiguous' if (type(_i) == np.float and np.isnan(_i)) else _i for _i in freq.index]
            freq = freq.sort_index(ascending = True)
        elif data_type == 'j_del_len':
            freq = self.indels.j_del_len.value_counts(dropna = False)
            freq.index = ['ambiguous' if (type(_i) == np.float and np.isnan(_i)) else _i for _i in freq.index]
            freq = freq.sort_index(ascending = True)
        elif data_type == 'vj_ins_len':
            freq = self.indels.vj_ins_len.value_counts(dropna = False)
            freq.index = ['ambiguous' if (type(_i) == np.float and np.isnan(_i)) else _i for _i in freq.index]
            freq = freq.sort_index(ascending = True)
        elif data_type == 'cdr3_prot_len':
            # set dropna as True to remove the 0-length CDR3
            freq = self.cdr3.prot_len.value_counts(dropna = True)
            freq.index = ['ambiguous' if (type(_i) == np.float and np.isnan(_i)) else _i for _i in freq.index]
            freq = freq.sort_index(ascending = True)
        elif data_type == 'cdr3_nucl_len':
            # set dropna as True to remove the 0-length CDR3
            freq = self.cdr3.nucl_len.value_counts(dropna = True)
            freq.index = ['ambiguous' if (type(_i) == np.float and np.isnan(_i)) else _i for _i in freq.index]
            freq = freq.sort_index(ascending = True)
        elif data_type == 'v':
            freq = self.vdj.v.value_counts(dropna = False)
        elif data_type == 'd':
            freq = self.vdj.d.value_counts(dropna = False)
            freq.index = ['ambiguous' if (type(_i) == np.float and np.isnan(_i)) else _i for _i in freq.index]
        elif data_type == 'j':
            freq = self.vdj.j.value_counts(dropna = False)
        if prob:
            freq = freq / freq.sum()

        return freq

    def __get_freq_vdj(self, data_type, prob):
        __sep = '__________'
        vdj_combinations = self.vdj.v.replace(np.nan, 'NA') + __sep + \
                           self.vdj.d.replace(np.nan, 'NA') + __sep + \
                           self.vdj.j.replace(np.nan, 'NA')
        freq = vdj_combinations.value_counts(dropna = False)

        if prob:
            freq = freq / freq.sum(axis = 2)

        vdj_v = []
        vdj_d = []
        vdj_j = []
        for vdj_combination in freq.index:
            vv, dd, jj = vdj_combination.split(__sep)
            vdj_v.append(vv)
            vdj_d.append(dd)
            vdj_j.append(jj)
        freq = pd.DataFrame({'v': vdj_v, 'd': vdj_d, 'j': vdj_j, 'frequency': freq.values},
                             columns = ['v', 'd', 'j', 'frequency']).replace('NA', np.nan)
        return freq

    def __get_est_vdj_rarefaction(self, data_type, func = 'mean'):
        rdata = None
        if self.div.rarefaction['vdj'] is not None:
            if func == 'raw':
                rdata = self.div.rarefaction['vdj']
            elif func == 'mean':
                rdata = self.div.rarefaction['vdj'].mean(axis = 1)
        return rdata

    def samplingresampling_study(self, data = None, n = 1000):
        """Sampling-resampling study.

        Args:
            data (str): ``vdj`` or ``cdr3``.

        Sampling-resamplig study.
        """

        if data is None or data == 'all':
            data = ['vdj']
        if not isinstance(data, list):
            data = [data]

        for data_i in data:
            if data_i == 'vdj':
                dat = self.vdj
                __sep = '__________'
                dat_for_study = dat.v.replace(np.nan, 'NaN') + __sep + \
                                dat.d.replace(np.nan, 'NaN') + __sep + \
                                dat.j.replace(np.nan, 'NaN')
                dat_for_study = pd.Series(dat_for_study)

            if data_i == 'cdr3':
                dat_for_study = pd.Series(self.cdr3.prot_seq)

            self.div.samplingresampling[data_i] = self.__samplingresampling_study(dat_for_study, n)

    def __samplingresampling_study(self, dat, n):
        population       = dat
        population_size  = len(dat)
        resampling_size_ratio = range(101)

        # get resampling sizes
        resampling_sizes = []
        for r in resampling_size_ratio:
            resampling_sizes.append(int(round(0.01 * r * population_size)))
        x = [0]   # total sampling sizes
        y = []    # the number of new sampled CDR3 sequence
        for i in range(n):
            y_try_n      = []
            population_n = list(population)
            sampled_n    = set([])

            for s in range(len(resampling_sizes) - 1):
                samplingsize = resampling_sizes[s + 1] - resampling_sizes[s]
                if len(population_n) < samplingsize * 1.5:
                    samplingsize = len(population_n)
                if len(population_n) <= samplingsize:
                    sampled_idx = range(0, len(population_n))
                else:
                    sampled_idx = random.sample(range(0, len(population_n)), samplingsize)
                s_left = []
                s_smpl = []
                for smplidx_in_population_n in range(len(population_n)):
                    if smplidx_in_population_n in sampled_idx:
                        s_smpl.append(population_n[smplidx_in_population_n])
                    else:
                        s_left.append(population_n[smplidx_in_population_n])
                sampled_items = s_smpl
                population_n  = s_left
                sampled_items = set(sampled_items)
                y_try_n.append(len(sampled_items.difference(sampled_n)))
                sampled_n = sampled_n.union(sampled_items)
                if i == 0:
                    x.append(x[s] + samplingsize)
            y.append(y_try_n)
        x.pop(0)
        y = pd.DataFrame(y, columns = resampling_size_ratio[1:]).T
        y.columns = ['try_' + str(n_try + 1) for n_try in range(n)]
        return y

    def rarefaction_study(self, data = None, n = 1000):
        """Perform rarefaction study for VDJ combination or CDR3 sequenece.

        Args:
            data (str): One of ``vdj`` or ``cdr3`` can be specified for diversity study.
            n (int): The number of performing of capture-recapture procedures.
        """

        if data is None or data == 'all':
            data = ['vdj']
        if not isinstance(data, list):
            data = [data]

        for data_i in data:
            if data_i == 'vdj':
                dat = self.vdj
                __sep = '__________'
                dat_for_study = dat.v.replace(np.nan, 'NaN') + __sep + \
                                dat.d.replace(np.nan, 'NaN') + __sep + \
                                dat.j.replace(np.nan, 'NaN')
                dat_for_study = pd.Series(dat_for_study)

            if data_i == 'cdr3':
                dat_for_study = pd.Series(self.cdr3.prot_seq)

            self.div.rarefaction[data_i] = self.__rarefaction_study(dat_for_study, n)

    def __rarefaction_study(self, dat, n):
        population_size      = len(dat)
        sampling_sizes       = self.__get_sampling_sizes(population_size)
        sampled_unique_items = [None] * n

        # rarefaction study
        for i in range(n):
            sampled_unique_items_i = []
            for s in sampling_sizes:
                sampled_id = random.sample(range(0, population_size), s)
                sampled_unique_items_i.append(dat[sampled_id].unique().shape[0])
            sampled_unique_items[i] = sampled_unique_items_i

        # set results into DataFrame class object
        sampled_unique_items = pd.DataFrame(sampled_unique_items, columns = sampling_sizes).T
        sampled_unique_items.columns = ['try_' + str(n_try + 1) for n_try in range(n)]

        return sampled_unique_items

    def __get_sampling_sizes(self, sample_size):
        sampling_sizes = []
        n_digits = int(math.log10(sample_size) + 1)
        # 1st sampling sizes
        for i in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]:
            sampling_sizes.append(i * 10 ** (n_digits - 1 - 2))
        # 2nd sampling sizes
        for j in [2, 4, 6, 8, 10]:
            sampling_sizes.append(j * 10 ** (n_digits - 1 - 1))
        # 3rd sampling sizes
        k = 2
        while max(sampling_sizes) < sample_size:
            if k * 10 ** (n_digits - 1) < sample_size:
                sampling_sizes.append(k * 10 ** (n_digits - 1))
            else:
                sampling_sizes.append(sample_size)
            k += 1
        return sampling_sizes

class PyDAIRStatsRecords:
    '''PyDAIRStatsRecords for storing the PyDAIR data.

    This class provides the list of PyDAIRStatsRecord and some function of list.
    '''

    def __init__(self):
        """PyDAIRStatsRecords class initialize method.

        """

        self.__records = []
        self.__index = 0

    def __len__(self):
        return len(self.__records)

    def len(self):
        """Retrive the number of PyDAIRStatsRecord class object.

        Retrive the number of PyDAIRStatsRecord objects, thus the number of samples.
        """

        return self.__len__()

    def __iter__(self):
        self.__index = 0
        return self

    def __next__(self):
        try:
            stats_record = self.__records[self.__index]
        except IndexError:
            raise StopIteration
        self.__index += 1
        return stats_record

    def next(self):
        """Return the next IgSeq object from the iterator of PyDAIRStatsRecords class objects.

        """

        return self.__next__()

    def append(self, stats_record):
        """Append PyDAIRStatsRecord into PyDAIRStatsRecords class object.

        """

        self.__records.append(stats_record)

    def get_record(self, i):
        """Retrive A PyDAIRStatsRecord.

        Args:
            i (int): The index of PyDAIRStatsRecord should be returned.

        Returns:
            PyDAIRStatsRecord: Return a PyDAIRStatsRecord object.
        """

        try:
           stats_record = self.__records[i]
        except IndexError:
            raise IndexError
        return stats_record

    def set_record(self, i, stats_record):
        """Update PyDAIRStatsRecord into PyDAIRStatsRecords class object.

        Args:
            i (int): The index of PyDAIRStatsRecord should be updated.
            stats_record: A PyDAIRStatsRecord that will used for updating.
        """

        self.__records[i] = stats_record

class PyDAIRStats:
    '''
    Statistics analysis of repertoire sequences.

    The class for storing the PyDAIR data file. Before use initialize this class,
    one should run create PyDAIR format files by other PyDAIR functions. This class
    stored many data of analyzed data.
    '''

    def __init__(self, pydair_file, pydair_id = None,
                 discard_ambiguous_D = False, productive_only = False):
        """PyDAIRStats class initialize method.

        Args:
            pydair_file (list): A list the contains multiple PYDAIR file path.
            pydair_id (list): A list of sample names.
            discard_ambiguous_D (bool): If true, discard ambiguous D before analysis.
            productive_only (bool): If true, analyze sequences with stop codons.
        """

        if pydair_id is None:
            pydair_id = []
            for i in range(len(pydair_file)):
                pydair_id.append('individual ' + str(i + 1))

        self.__pydair_file   = pydair_file
        self.__pydair_id     = pydair_id
        self.__discard_ambiguous_D = discard_ambiguous_D
        self.__productive_only = productive_only
        self.samples   = None

        # parse PyDAIR files
        self.__parse_pydair_files()

    def __parse_pydair_files(self):
        self.samples  = PyDAIRStatsRecords()
        for i in range(len(self.__pydair_file)):
            v = []
            d = []
            j = []
            orf = []
            cdr3_prot_seq = []
            cdr3_nucl_seq = []
            v_del = []
            j_del = []
            vj_ins = []

            pydair_fh = PyDAIRIO(self.__pydair_file[i], 'r')
            for igseq in pydair_fh.parse():
                v.append(igseq.v.sbjct.name)
                d.append(igseq.d.sbjct.name)
                j.append(igseq.j.sbjct.name)
                orf.append(igseq.query.orf)

                cdr3_data = igseq.get_cdr3_data()
                if cdr3_data.nucl_seq is not None and '*' in cdr3_data.prot_seq:
                    cdr3_data.nucl_seq = None
                    cdr3_data.prot_seq = None
                cdr3_prot_seq.append(cdr3_data.prot_seq)
                cdr3_nucl_seq.append(cdr3_data.nucl_seq)

                v_del.append(igseq.indels.v_deletion)
                j_del.append(igseq.indels.j_deletion)
                vj_ins.append(igseq.indels.vj_insertion)

            sample_record = PyDAIRStatsRecord(self.__pydair_id[i], v, d, j, orf,
                                              cdr3_nucl_seq, cdr3_prot_seq,
                                              v_del, j_del, vj_ins,
                                              self.__discard_ambiguous_D, self.__productive_only)
            self.samples.append(sample_record)
            pydair_fh.close()

    def rarefaction_study(self, data = None, n = 1000):
        """Rarefaction study for estimating the number of VDJ combination.

        Args:
            data (str): 'vdj' or 'cdr3'.

        Rarefaction study for estimating the number of VDJ combination for each sample.
        """

        for bsample in self.samples:
            bsample.rarefaction_study(data, n)

    def get_summary(self, data_type, prob = False):
        '''Get the frequence of usage with data frame class object.

        Args:
            data_type (str): A string to specify data type. ``v``,
                             ``d``, ``j``, ``vdj``, ``cdr3_nucl_len``,
                             ``cdr3_prot_len``, ``v_del_len``, ``j_del_len``,
                             and ``vj_ins_len`` are supported.
            freq (bool): If `True`, return the frequences of counts.
            prob (bool): If prob is `True`, return the probability, and omit `freq`.

        Returns:
            A Pandas DataFrame class object.

        Get the frequence of usages of the ``data_type`` as DataFrame class object.
        '''

        s = None

        if data_type in ['v', 'd', 'j', 'cdr3_prot_len', 'cdr3_nucl_len', 'v_del_len', 'j_del_len', 'vj_ins_len']:
            s = self.__get_freq(data_type, prob)
        elif data_type == 'vdj':
            s = self.__get_freq_vdj(data_type, prob)
        elif data_type == 'vdj_rarefaction':
            s = self.__get_est_vdj_rarefaction(data_type)
        else:
            raise ValueError(data_type + ' is not supported by get_summary method.')

        return s

    def __get_freq(self, data_type, prob = False):
        sample_freqs = []
        sample_names = []
        for sample in self.samples:
            sample_freqs.append(sample.get_summary(data_type, prob = prob))
            sample_names.append(sample.name)
        freq_dataframe = pd.concat(sample_freqs, axis = 1)
        freq_dataframe.columns = sample_names

        if prob:
            freq_dataframe = freq_dataframe / freq_dataframe.sum(axis = 1)

        freq_dataframe.columns = sample_names

        if data_type in ['v', 'd', 'j']:
            freq_dataframe = freq_dataframe.ix[freq_dataframe.mean(axis = 1).sort_values(ascending = False).index]
        else:
            freq_dataframe = freq_dataframe.set_index([[int(dx) for dx in freq_dataframe.index.values]])
        return freq_dataframe

    def __get_freq_vdj(self, data_type, prob = False):
        __sep = '__________'
        sample_freqs = []
        sample_names = []
        for sample in self.samples:
            freq = sample.get_summary(data_type, prob = prob)
            freqval = pd.Series(freq.frequency)
            freqval.index = freq.v.replace(np.nan, 'NA') + __sep + \
                            freq.d.replace(np.nan, 'NA') + __sep + \
                            freq.j.replace(np.nan, 'NA')
            sample_freqs.append(freqval)
            sample_names.append(sample.name)
        freq_dataframe = pd.concat(sample_freqs, axis = 1)
        freq_dataframe.columns = sample_names

        vdj_v = []
        vdj_d = []
        vdj_j = []
        for vdj_combination in freq_dataframe.index:
            vv, dd, jj = vdj_combination.split(__sep)
            vdj_v.append(vv)
            vdj_d.append(dd)
            vdj_j.append(jj)

        freq_vdjcmbn = pd.concat([pd.Series(vdj_v), pd.Series(vdj_d), pd.Series(vdj_j)], axis = 1).replace('NA', np.nan)
        freq_vdjcmbn.columns = ['V', 'D', 'J']

        freq_vdjcmbn.reset_index(drop = True, inplace = True)
        freq_dataframe.reset_index(drop = True, inplace = True)

        freq = pd.concat([freq_vdjcmbn, freq_dataframe], axis = 1)
        return freq

    def __get_est_vdj_rarefaction(self, data_type):
        rdata = []
        sample_names = []
        for sample in self.samples:
            sample_rdata = sample.get_summary('vdj_rarefaction', 'mean')
            sample_names.append(sample.name)
            if sample_rdata is not None:
                rdata.append(sample_rdata)
        if len(rdata) > 0:
            rdata = pd.concat(rdata, axis = 1)
            rdata.columns = sample_names
        else:
            rdata = None

        return rdata


"""
Routines for getting temperature data for starcryo cryostat
Currently assumes SRS logs are in ~heather/SRS
"""

import glob
import os
import time
import bisect
import numpy as np
import netCDF4
import kid_readout.utils.parse_srs

temperature_log_file_dir = '/home/heather/SRS/'

_filecache = {}
epochs = []
temperature_log_filenames = []
temperature_log_filenames_last_updated = 0

def get_temperature_log_file_list():
    # First the original logs
    log_filenames = glob.glob(os.path.join(temperature_log_file_dir,'2*.txt'))
    log_filenames.sort()
    if len(log_filenames) == 0:
        raise Exception("Could not find any temperature log files in %s" % temperature_log_file_dir)
    epochs = []
    for filename in log_filenames:
        base,fname = os.path.split(filename)
        epochs.append(time.mktime(time.strptime(fname,'%Y%m%d-%H%M%S.txt')))

    epochs = np.array(epochs)
    return epochs,log_filenames

def refresh_temperature_log_file_list():
    global epochs
    global temperature_log_filenames
    global temperature_log_filenames_last_updated
    if not temperature_log_filenames or time.time() - temperature_log_filenames_last_updated > 20*3600:
        epochs,temperature_log_filenames = get_temperature_log_file_list()
        temperature_log_filenames_last_updated = time.time()
        print "updated file list"

refresh_temperature_log_file_list()

def get_temperatures_at(t):
    refresh_temperature_log_file_list()
    global epochs
    global temperature_log_filenames
    global _filecache

    if np.isscalar(t):
        start_time = t
        end_time = t
    else:
        start_time = t.min()
        end_time = t.max()
    idx = bisect.bisect_right(epochs,start_time)
    idx = idx - 1
    if idx < 0:
        idx = 0
    filename = temperature_log_filenames[idx]

    if _filecache.has_key(filename) and (time.time()-end_time) > 20*60: #if we're looking for data from more than 20 minutes ago, look in the cache
        datetime_timestamp,data = _filecache[filename]
    else:
        datetime_timestamp,data = kid_readout.utils.parse_srs.get_load_log(filename)
        _filecache[filename] = (datetime_timestamp,data)

    times = data[:,0]
    if end_time > times[-1]:
        print "Warning: requested times may span more than one log file, so results may not be as intended"
        print "log file is: %s, last requested time is %s" % (filename, time.ctime(end_time))
    primary_package_temperature = np.interp(t,times,data[:,-4])
    secondary_package_temperature = np.interp(t,times,data[:,-2])
    primary_load_temperature = np.interp(t,times,data[:,1])
    secondary_load_temperature = np.interp(t,times,data[:,2])
    return primary_package_temperature, secondary_package_temperature, primary_load_temperature, secondary_load_temperature


#!/usr/local/bin/python3
# -*- coding: utf-8 -*-
import sys
import os
# print(os.listdir())
# sys.path.append('/models/SingleHumanParser')
import os
import argparse
import logging
import numpy as np
from PIL import Image
import matplotlib
import matplotlib.pyplot as plt
import torch

import torch.nn as nn
from torchvision import transforms

#from net.pspnet import PSPNet
from net.pspnet import PSPNet
import cv2 as cv
models = {
    'squeezenet': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='squeezenet'),
    'densenet': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=1024, deep_features_size=512, backend='densenet'),
    'resnet18': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet18'),
    'resnet34': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet34'),
    'resnet50': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet50'),
    'resnet101': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet101'),
    'resnet152': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet152')
}

parser = argparse.ArgumentParser(description="Pyramid Scene Parsing Network")
parser.add_argument('--image_path', type=str, default='', help='Path to image')
parser.add_argument('--models-path', type=str, default='./checkpoints', help='Path for storing model snapshots')
parser.add_argument('--backend', type=str, default='densenet', help='Feature extractor')
parser.add_argument('--num-classes', type=int, default=20, help="Number of classes.")
args = parser.parse_args()
args.models_path = './models/SingleHumanParser/checkpoints'
# image_folder = '/Users/qingyunw/Downloads/Origin_Data/train'
#
# image_names = sorted(os.listdir(image_folder))
# image_path = os.path.join(image_folder, image_names[1])
# parser.image_path = image_path
# img = Image.open(image_path)

def build_network(snapshot, backend):
    epoch = 0
    backend = backend.lower()
    net = models[backend]()
    net = nn.DataParallel(net)
    if snapshot is not None:
        _, epoch = os.path.basename(snapshot).split('_')
        if not epoch == 'last':
            epoch = int(epoch)
        #net.load_state_dict(torch.load(snapshot, map_location=torch.device('cpu')))
        net.load_state_dict(torch.load(snapshot))
        logging.info("Snapshot for epoch {} loaded from {}".format(epoch, snapshot))
    if torch.cuda.is_available():
        net = net.cuda()
    return net, epoch

def get_transform():
    transform_image_list = [
        #transforms.Resize((192, 256), 3),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
    return transforms.Compose(transform_image_list)

def show_image(img, pred, name):
    fig, axes = plt.subplots(1, 2)
    ax0, ax1 = axes
    ax0.get_xaxis().set_ticks([])
    ax0.get_yaxis().set_ticks([])
    ax1.get_xaxis().set_ticks([])
    ax1.get_yaxis().set_ticks([])

    classes = np.array(('Background',  # always index 0
                        'Hat', 'Hair', 'Glove', 'Sunglasses',
                        'UpperClothes', 'Dress', 'Coat', 'Socks',
                        'Pants', 'Jumpsuits', 'Scarf', 'Skirt',
                        'Face', 'Left-arm', 'Right-arm', 'Left-leg',
                        'Right-leg', 'Left-shoe', 'Right-shoe',))
    colormap = [(0, 0, 0),
                (1, 0.25, 0), (0, 0.25, 0), (0.5, 0, 0.25), (1, 1, 1),
                (1, 0.75, 0), (0, 0, 0.5), (0.5, 0.25, 0), (0.75, 0, 0.25),
                (1, 0, 0.25), (0, 0.5, 0), (0.5, 0.5, 0), (0.25, 0, 0.5),
                (1, 0, 0.75), (0, 0.5, 0.5), (0.25, 0.5, 0.5), (1, 0, 0),
                (1, 0.25, 0), (0, 0.75, 0), (0.5, 0.75, 0), ]
    cmap = matplotlib.colors.ListedColormap(colormap)
    bounds = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
    norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)

    h, w, _ = pred.shape

    def denormalize(img, mean, std):
        c, _, _ = img.shape
        for idx in range(c):
            img[idx, :, :] = img[idx, :, :] * std[idx] + mean[idx]
        return img

    img = denormalize(img.cpu().numpy(), [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    img = img.transpose(1, 2, 0).reshape((h, w, 3))
    pred = pred.reshape((h, w))

    # show image
    ax0.set_title('img')
    ax0.imshow(img)
    ax1.set_title('pred')
    mappable = ax1.imshow(pred, cmap=cmap, norm=norm)
    # colorbar legend
    cbar = plt.colorbar(mappable, ax=axes, shrink=0.7, )
    cbar.ax.get_yaxis().set_ticks([])
    for j, lab in enumerate(classes):
        cbar.ax.text(2.3, (j + 0.45) / 20.0, lab, ha='left', va='center', )

    plt.savefig(fname="./"  + name)
    print('result saved to ...')
    plt.show()

def class_change(n):
    # classes = np.array(('Background',  # always index 0
    #                     'Hat', 'Hair', 'Glove', 'Sunglasses',
    #                     'UpperClothes', 'Dress', 'Coat', 'Socks',
    #                     'Pants', 'Jumpsuits', 'Scarf', 'Skirt',
    #                     'Face', 'Left-arm', 'Right-arm', 'Left-leg',
    #                     'Right-leg', 'Left-shoe', 'Right-shoe',))
        # 0 -> Background
        # 1 -> Hair
        # 4 -> Upclothes
        # 5 -> Left-shoe
        # 6 -> Right-shoe
        # 8 -> Pants
        # 9 -> Left_leg
        # 10 -> Right_leg
        # 11 -> Left_arm
        # 12 -> Face
        # 13 -> Right_arm
    table = {
      0:0,
      1:1,
      2:1,
      3:0,
      4:12,
      5:4,
      6:4,
      7:4,
      8:0,
      9:8,
      10:4,
      11:4,
      12:4,
      13:12,
      14:11,
      15:13,
      16:9,
      17:10,
      18:5,
      19:6
    }
    return table[n]

def main():
    # --------------- model --------------- #
    snapshot = os.path.join(args.models_path, args.backend, 'PSPNet_last')
    net, starting_epoch = build_network(snapshot, args.backend)
    net.eval()

    # ------------ load image ------------ #
    data_transform = get_transform()
    img = Image.open(image_path)

    #img = Image.open(open(args.image_path, 'rb'))
    img = data_transform(img)
    if torch.cuda.is_available():
        img = img.cuda()

    # --------------- inference --------------- #

    with torch.no_grad():
        pred, _ = net(img.unsqueeze(dim=0))
        pred = pred.squeeze(dim=0)
        pred = pred.cpu().numpy().transpose(1, 2, 0)
        #pred = np.asarray(np.argmax(pred, axis=2), dtype=np.uint8).reshape((256, 256, 1))
        pred = np.asarray(np.argmax(pred, axis=2), dtype=np.uint8).reshape((256, 192, 1))
        vfunc = np.vectorize(class_change)
        pred_change = vfunc(pred)
        # status = cv.imwrite('/content/gdrive/My Drive/Insight_Project/DeepFashion_Try_On/Data_preprocessing/test_label/0000_0.png',pred_change)
        status = cv.imwrite('./sample/0000_1.png',pred_change)

        show_image(img, pred, '0_')
        show_image(img, pred_change, '1_')

def get_parser(person_path, parser_path):
    # --------------- load model --------------- #
    snapshot = os.path.join(args.models_path, args.backend, 'PSPNet_last')
    net, starting_epoch = build_network(snapshot, args.backend)
    net.eval()

    # ------------ load image ------------ #
    data_transform = get_transform()
    img = Image.open(person_path)
    img = data_transform(img)
    if torch.cuda.is_available():
        img = img.cuda()

    # --------------- inference --------------- #

    with torch.no_grad():
        pred, _ = net(img.unsqueeze(dim=0))
        pred = pred.squeeze(dim=0)
        pred = pred.cpu().numpy().transpose(1, 2, 0)
        #pred = np.asarray(np.argmax(pred, axis=2), dtype=np.uint8).reshape((256, 256, 1))
        pred = np.asarray(np.argmax(pred, axis=2), dtype=np.uint8).reshape((256, 192, 1))
        vfunc = np.vectorize(class_change)
        pred_change = vfunc(pred)
        # status = cv.imwrite('/content/gdrive/My Drive/Insight_Project/DeepFashion_Try_On/Data_preprocessing/test_label/0000_0.png',pred_change)
        status = cv.imwrite(parser_path, pred_change)
        return parser_path
        #
        # colormap = [(0, 0, 0),
        #             (1, 0.25, 0), (0, 0.25, 0), (0.5, 0, 0.25), (1, 1, 1),
        #             (1, 0.75, 0), (0, 0, 0.5), (0.5, 0.25, 0), (0.75, 0, 0.25),
        #             (1, 0, 0.25), (0, 0.5, 0), (0.5, 0.5, 0), (0.25, 0, 0.5),
        #             (1, 0, 0.75), (0, 0.5, 0.5), (0.25, 0.5, 0.5), (1, 0, 0),
        #             (1, 0.25, 0), (0, 0.75, 0), (0.5, 0.75, 0), ]
        # cmap = matplotlib.colors.ListedColormap(colormap)
        # bounds = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
        # norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)
        # h, w, _ = pred_change.shape
        #
        # pred_change = pred_change.reshape((h, w))
        #
        # # show image
        # plt.imshow(pred_change, cmap=cmap, norm=norm)
        # plt.savefig(fname="./" + name)
        # print('result saved to ...')
    #return Image.open('./sample/0000_1.png')

if __name__ == '__main__':
    main()


"""
pyrad.io.trajectory
===================

Trajectory class implementation for reading trajectory file.
Converting to different coordinate systems.

.. autosummary::
    :toctree: generated/

    Trajectory
    _Radar_Trajectory

"""

import sys
import re
import datetime
import locale
from warnings import warn
from copy import deepcopy

import numpy as np

import pyart

from ..io.read_data_sensor import read_lightning, read_trt_traj_data
from ..io.read_data_sensor import read_trt_thundertracking_traj_data

class Trajectory(object):
    """
    A class for reading and handling trajectory data from a file.

    Attributes
    ----------
    filename : str
        Path and name of the trajectory definition file
    starttime : datetime
        Start time of trajectory processing.
    endtime : datetime
        End time of trajectory processing.
    trajtype : str
    Type of trajectory. Can be 'plane' or 'lightning'
    time_vector : Array of datetime objects
        Array containing the trajectory time samples
    wgs84_lat_deg : Array of floats
        WGS84 latitude samples in radian
    wgs84_lon_deg : Array of floats
        WGS84 longitude samples in radian
    wgs84_alt_m : Array of floats
        WGS84 altitude samples in m
    nsamples : int
    Number of samples in the trajectory
    _swiss_grid_done : Bool
        Indicates that convertion to Swiss coordinates has been performed
    swiss_chy, swiss_chx, swiss_chh : Array of floats
        Swiss coordinates in m
    radar_list : list
          List of radars for which trajectories are going to be computed
    flashnr : int
          For 'lightning' only. Number of flash for which trajectory data
          is going to be computed. If 0 all all flashes are going to be
          considered.
    time_in_flash : array of floats
          For 'lightning' only. Time within flash (sec)
    flashnr_vec : array of ints
          For 'lightning' only. Flash number of each data sample
    dBm : array of floats
          For 'lightning' only. Lightning power (dBm)

    Methods:
    --------
    add_radar : Add a radar
    calculate_velocities : Computes velocities
    get_start_time : Return time of first trajectory sample
    get_end_time : Return time of last trajectory sample
    get_samples_in_period : Get indices of samples within period
    _convert_traj_to_swissgrid : convert data from WGS84 to Swiss coordinates
    _read_traj : Read plane trajectory from file
    _read_traj_trt : Read TRT trajectory from file
    _read_traj_lightning : Read lightning trajectory from file
    _get_total_seconds : Get the total time of the trajectory in seconds

    """

    def __init__(self, filename, starttime=None, endtime=None,
                 trajtype='plane', flashnr=0):
        """
        Initalize the object.

        Parameters
        ----------
        filename : str
            Filename containing the trajectory samples.
        starttime : datetime
            Start time of trajectory processing. If not given, use
            the time of the first trajectory sample.
        endtime : datetime
            End time of trajectory processing. If not given, use
            the time of the last trajectory sample.
        trajtype : str
            type of trajectory. Can be plane or lightning
        flashnr : int
            If type of trajectory is lightning, the flash number to check the
            trajectory. 0 means all flash numbers included
        """

        self.filename = filename
        self.starttime = starttime
        self.endtime = endtime
        self.trajtype = trajtype

        self.time_vector = np.array([], dtype=datetime.datetime)
        self.wgs84_lat_deg = np.array([], dtype=float)
        self.wgs84_lon_deg = np.array([], dtype=float)
        self.wgs84_alt_m = np.array([], dtype=float)
        self.nsamples = None

        self._swiss_grid_done = False
        self.swiss_chy = None
        self.swiss_chx = None
        self.swiss_chh = None

        self.radar_list = []

        if self.trajtype == 'lightning':
            self.flashnr = flashnr
            self.time_in_flash = np.array([], dtype=datetime.datetime)
            self.dBm = np.array([], dtype=float)
            self.flashnr_vec = np.array([], dtype=float)
        elif self.trajtype == 'trt':
            self.cell_contour = np.array([])

        try:
            if self.trajtype == 'lightning':
                self._read_traj_lightning(flashnr)
            elif self.trajtype == 'trt':
                self._read_traj_trt()
            else:
                self._read_traj()

        except Exception as ee:
            warn(str(ee))
            raise Exception(
                "ERROR: Could not load trajectory data from file '" +
                filename+"' into Trajectory object")

    def add_radar(self, radar):
        """
        Add the coordinates (WGS84 longitude, latitude and non WGS84 altitude)
        of a radar to the radar_list.

        Parameters
        ----------
        radar : pyart radar object
            containing the radar coordinates

        Return
        ------
        Radar object

        """

        # Check if radar location is already in the radar list
        for rad in self.radar_list:
            if (rad.location_is_equal(radar.latitude['data'][0],
                                      radar.longitude['data'][0],
                                      radar.altitude['data'][0])):
                # warn("WARNING: Tried to add the same radar twice to the"
                #      " radar list")
                return rad

        rad = _Radar_Trajectory(radar.latitude['data'][0],
                                radar.longitude['data'][0],
                                radar.altitude['data'][0])
        self.radar_list.append(rad)

        # Convert trajectory WGS84 points to polara radar coordinates
        rad.convert_radpos_to_swissgrid()

        self._convert_traj_to_swissgrid()

        # Note: Earth curvature not considered yet!

        (rvec, azvec, elvec) = pyart.core.cartesian_to_antenna(
            self.swiss_chy-rad.ch_y, self.swiss_chx-rad.ch_x,
            self.swiss_chh-rad.ch_alt)

        rad.assign_trajectory(elvec, azvec, rvec)

        return rad

    def calculate_velocities(self, radar):
        """
        Calculate velocities.

        """

        if not radar.traj_assigned:
            raise Exception("ERROR: No trajectory assigned to radar object")

        dt_secs = np.vectorize(self._get_total_seconds)
        dt = dt_secs(self.time_vector[2:] - self.time_vector[:-2])

        v_r = np.empty(self.nsamples, dtype=float)
        v_r[0] = v_r[-1] = np.nan
        v_r[1:-1] = (radar.range_vec[2:] - radar.range_vec[:-2]) / dt

        v_az = np.empty(self.nsamples, dtype=float)
        v_az[0] = v_az[-1] = np.nan
        daz = radar.azimuth_vec[2:] - radar.azimuth_vec[:-2]
        daz[daz > 180.] -= 360.
        daz[daz < -180.] += 360.
        v_az[1:-1] = daz / dt

        v_el = np.empty(self.nsamples, dtype=float)
        v_el[0] = v_el[-1] = np.nan
        v_el[1:-1] = (radar.elevation_vec[2:] - radar.elevation_vec[:-2]) / dt

        v_abs = np.empty(self.nsamples, dtype=float)
        v_abs[0] = v_abs[-1] = np.nan
        self._convert_traj_to_swissgrid()
        dx = self.swiss_chy[2:] - self.swiss_chy[:-2]
        dy = self.swiss_chx[2:] - self.swiss_chx[:-2]
        dz = self.swiss_chh[2:] - self.swiss_chh[:-2]
        v_abs[1:-1] = np.sqrt(dx**2 + dy**2 + dz**2) / dt

        radar.assign_velocity_vecs(v_abs, v_r, v_el, v_az)

    def get_samples_in_period(self, start=None, end=None):
        """"
        Get indices of samples of the trajectory within given time
        period.
        """

        if ((start is None) and (end is None)):
            raise Exception("ERROR: Either start or end must be defined")
        elif start is None:
            return np.where(self.time_vector < end)
        elif end is None:
            return np.where(self.time_vector >= start)
        else:
            return np.where((self.time_vector >= start) &
                            (self.time_vector < end))

    def get_start_time(self):
        """
        Get time of first trajectory sample.

        Return
        ------
        datetime object
        """

        return self.time_vector[0]

    def get_end_time(self):
        """
        Get time of last trajectory sample.

        Return
        ------
        datetime object
        """

        return self.time_vector[-1]

    def _convert_traj_to_swissgrid(self):
        """
        Convert trajectory samples from WGS84 to Swiss CH1903 coordinates
        """

        if self._swiss_grid_done:
            return

        (self.swiss_chy, self.swiss_chx, self.swiss_chh) = \
            pyart.core.wgs84_to_swissCH1903(
                self.wgs84_lon_deg, self.wgs84_lat_deg, self.wgs84_alt_m)

        self._swiss_grid_done = True

    def _read_traj(self):
        """
        Read trajectory from file

        File format
        -----------
        Comment symbol: '#'
        Columns:
        1. Day (UTC) format: DD-MMM-YYYY
        2. Seconds since midnight (UTC)
        3. WGS84 latitude [radians]
        4. WGS84 longitude [radians]
        5. WGS84 altitude [m]

        """

        # check if the file can be read
        try:
            tfile = open(self.filename, "r")
        except:
            raise Exception("ERROR: Could not find|open trajectory file '" +
                            self.filename+"'")

        repat = re.compile("(\d+\-[A-Za-z]+\-\d+)\s+([\d\.]+)\s+"
                           "([\-\d\.]+)\s+([\-\d\.]+)\s+([\-\d\.]+)")

        try:
            loc_set = False
            loc = locale.getlocale()  # get current locale
            if loc[0] != 'en_US':
                try:
                    locale.setlocale(locale.LC_ALL, ('en_US', 'UTF-8'))
                except Exception as ee:
                    raise Exception("ERROR: Cannot set local 'en_US': %s"
                                    % str(ee))
                loc_set = True

            recording_started = True
            if self.starttime is not None:
                recording_started = False
            recording_check_stop = False
            if self.endtime is not None:
                recording_check_stop = True

            for line in tfile:
                line = line.strip()

                if not line:
                    continue

                # ignore comments
                if line.startswith('#'):
                    continue

                line = line.partition('#')[0]  # Remove comments
                line = line.strip()

                mm = repat.match(line)
                if not mm:
                    print("WARNING: Format error in trajectory file '%s'"
                          " on line '%s'" % (self.filename, line),
                          file=sys.stderr)
                    continue

                # Get time stamp
                try:
                    sday = datetime.datetime.strptime(mm.group(1), "%d-%b-%Y")
                except Exception as ee:
                    print(datetime.datetime.utcnow().strftime("%d-%b-%Y"))
                    raise Exception("ERROR: Format error in traj file '%s' "
                                    "on line '%s' (%s)"
                                    % (self.filename, line, str(ee)))

                sday += datetime.timedelta(seconds=float(mm.group(2)))

                if not recording_started:
                    if sday < self.starttime:
                        continue
                    else:
                        recording_started = True

                if recording_check_stop:
                    if sday > self.endtime:
                        break

                self.time_vector = np.append(self.time_vector, [sday])

                self.wgs84_lat_deg = np.append(
                    self.wgs84_lat_deg, [float(mm.group(3)) * 180. / np.pi])
                self.wgs84_lon_deg = np.append(
                    self.wgs84_lon_deg, [float(mm.group(4)) * 180. / np.pi])
                self.wgs84_alt_m = np.append(
                    self.wgs84_alt_m, [float(mm.group(5))])
        except:
            raise
        finally:
            tfile.close()
            if loc_set:
                locale.setlocale(locale.LC_ALL, loc)  # restore saved locale

        self.nsamples = len(self.time_vector)

    def _read_traj_lightning(self, flashnr=0):
        """
        Read trajectory from lightning file

        File format
        -----------
        Columns:
        1. flashnr (0=noise)
        2. Seconds since midnight (UTC)
        2. Time in flash (s)
        3. WGS84 latitude [deg]
        4. WGS84 longitude [deg]
        5. WGS84 altitude [m]
        6. Power [dBm]

        Parameters
        ----------
        flashnr : int
            the flash number to keep. If 0 data from all flashes will be kept

        """
        flashnr_vec, time, time_in_flash, lat, lon, alt, dBm = read_lightning(
            self.filename)

        if flashnr_vec is None:
            raise Exception("ERROR: Could not find|open trajectory file '" +
                            self.filename+"'")

        recording_started = True
        if self.starttime is not None:
            recording_started = False
        recording_check_stop = False
        if self.endtime is not None:
            recording_check_stop = True

        if flashnr > 0:
            flashnr_vec_aux = deepcopy(flashnr_vec)
            flashnr_vec = flashnr_vec[flashnr_vec_aux == flashnr]
            time = time[flashnr_vec_aux == flashnr]
            time_in_flash = time_in_flash[flashnr_vec_aux == flashnr]
            lat = lat[flashnr_vec_aux == flashnr]
            lon = lon[flashnr_vec_aux == flashnr]
            alt = alt[flashnr_vec_aux == flashnr]
            dBm = dBm[flashnr_vec_aux == flashnr]

        for i, dBm_val in enumerate(dBm):

            if not recording_started:
                if time[i] < self.starttime:
                    continue
                else:
                    recording_started = True

            if recording_check_stop:
                if time[i] > self.endtime:
                    break

            self.flashnr_vec = np.append(self.flashnr_vec, [flashnr_vec[i]])
            self.time_vector = np.append(self.time_vector, [time[i]])
            self.time_in_flash = np.append(
                self.time_in_flash, [time_in_flash[i]])

            self.wgs84_lat_deg = np.append(self.wgs84_lat_deg, [lat[i]])
            self.wgs84_lon_deg = np.append(self.wgs84_lon_deg, [lon[i]])
            self.wgs84_alt_m = np.append(self.wgs84_alt_m, [alt[i]])

            self.dBm = np.append(self.dBm, [dBm_val])

        self.nsamples = len(self.time_vector)

    def _read_traj_trt(self):
        """
        Read trajectory from TRT file

        File format
        -----------
        Columns:
        1. traj_ID
        2. yyyymmddHHMM (UTC)
        2. lon (deg)
        3. lat [deg]
        4. ell_L
        5. ell_S
        6. ell_or
        7. area
        8. vel_x
        9. vel_y
        10. det
        11. RANKr
        12. CG_n
        13. CG_p
        14. CG
        15. CG_percent_p
        16. ET45
        17. ET45m
        18. ET15
        19. ET15m
        20. VIL
        21. maxH
        22. maxHm
        23. POH
        24. RANK
        25. Dvel_x
        26. Dvel_y
        27. cell_contours

        Parameters
        ----------

        """
        if '_tt.trt' in self.filename:
            (traj_ID, _, yyyymmddHHMM, _, _, _, lon, lat, _, _, _, _, _, _, _,
             _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _,
             cell_contours) = read_trt_thundertracking_traj_data(
                 self.filename)

            if traj_ID is None:
                raise Exception(
                    "ERROR: Could not find|open trajectory file '" +
                    self.filename+"'")

            valid = np.logical_not(np.ma.getmaskarray(yyyymmddHHMM))
            yyyymmddHHMM = yyyymmddHHMM[valid]
            traj_ID = traj_ID[valid]
            lon = lon[valid]
            lat = lat[valid]
            cell_contours = cell_contours[valid]

            if traj_ID.size == 0:
                raise Exception(
                    "ERROR: No valid data in trajectory file '" +
                    self.filename+"'")

        else:
            (traj_ID, yyyymmddHHMM, lon, lat, _, _, _, _, _, _, _, _, _, _, _,
             _, _, _, _, _, _, _, _, _, _, _, _, cell_contours) = (
                 read_trt_traj_data(self.filename))

            if traj_ID is None:
                raise Exception(
                    "ERROR: Could not find|open trajectory file '" +
                    self.filename+"'")

        recording_started = True
        if self.starttime is not None:
            recording_started = False
        recording_check_stop = False
        if self.endtime is not None:
            recording_check_stop = True

        for i, cell_contour in enumerate(cell_contours):
            if not recording_started:
                if yyyymmddHHMM[i] < self.starttime:
                    continue
                else:
                    recording_started = True

            if recording_check_stop:
                if yyyymmddHHMM[i] > self.endtime:
                    break

            self.time_vector = np.append(self.time_vector, [yyyymmddHHMM[i]])

            self.wgs84_lat_deg = np.append(self.wgs84_lat_deg, [lat[i]])
            self.wgs84_lon_deg = np.append(self.wgs84_lon_deg, [lon[i]])
            self.wgs84_alt_m = np.append(self.wgs84_alt_m, 0.)

            self.cell_contour = np.append(self.cell_contour, [cell_contour])

        self.nsamples = len(self.time_vector)

    def _get_total_seconds(self, x):
        """ Return total seconds of timedelta object"""
        return x.total_seconds()

class _Radar_Trajectory:
    """
    A class for holding the trajectory data assigned to a radar.

    Attributes
    ----------
    latitude : float
       WGS84 radar latitude [deg]
    longitude : float
       WGS84 radar longitude [deg]
    altitude : float
       radar altitude [m] (non WGS84)
    ch_y, ch_x, ch_alt : float
       radar coordinates in swiss CH1903 coordinates
    elevation_vec : float list
       Elevation values of the trajectory samples
    azimuth_vec : float list
       Azimuth values of the trajectory samples
    range_vec : float list
       Range values of the trajectory samples
    v_abs, v_r, v_el, v_az : array-like
       Velocity vectors of the absolute [m/s], radial [m/s], elevation [deg/s]
       and azimuth [deg/s] velocities

    Methods:
    --------
    location_is_equal
    assign_trajectory
    convert_radpos_to_swissgrid
    assign_velocity_vecs
    """

    def __init__(self, lat, lon, alt):
        """
        Initalize the object.

        Parameters
        ----------
        lat, lon , alt : radar location coordinates
        nsamps : number of samples
        """

        # radar position in lat, lon
        self.latitude = lat
        self.longitude = lon
        self.altitude = alt

        # radar position in swiss CH1903 coordinates
        self._swiss_coords_done = False
        self.ch_y = None
        self.ch_x = None
        self.ch_alt = None

        # trajectory in polar radar coordinates
        self.traj_assigned = False
        self.elevation_vec = None
        self.azimuth_vec = None
        self.range_vec = None

        # Velocity vectors
        self._velocity_vecs_assigned = False
        self.v_abs = None
        self.v_r = None
        self.v_el = None
        self.v_az = None

    def location_is_equal(self, lat, lon, alt):
        """
        Check if the given coordinates are the same.

        Parameters
        ----------
        lat, lon , alt : radar location coordinates

        Return
        ------
        True if the radar location is equal, False otherwise
        """

        lat_tol = 0.002  # [deg]
        lon_tol = 0.002  # [deg]
        alt_tol = 2.0    # [m]

        if ((np.abs(self.latitude-lat) > lat_tol) or
                (np.abs(self.longitude-lon) > lon_tol) or
                (np.abs(self.altitude-alt) > alt_tol)):
            return False

        return True

    def assign_trajectory(self, el, az, rr):
        """
        Assign a trajectory to the radar in polar radar
        coordinates.

        Parameters
        ----------
        el, az, rr : array-like
           elevation, azimuth and range vector
        """

        if self.traj_assigned:
            warn("WARNING: Trajectory already assigned")
            return

        self.elevation_vec = el
        self.azimuth_vec = az
        self.range_vec = rr
        self.traj_assigned = True

    def convert_radpos_to_swissgrid(self):
        """
        Convert the radar location (in WGS84 coordinates) to
        swiss CH1903 coordinates.

        """

        if self._swiss_coords_done:
            return

        (self.ch_y, self.ch_x, self.ch_alt) = \
            pyart.core.wgs84_to_swissCH1903(self.longitude,
                                            self.latitude,
                                            self.altitude,
                                            no_altitude_transform=True)
        self._swiss_coords_done = True

    def assign_velocity_vecs(self, v_abs, v_r, v_el, v_az):
        """
        Assign velocity vectors to the radar.

        """

        if self._velocity_vecs_assigned:
            warn("WARNING: Trajectory velocity vectors already assigned")
            return

        self.v_abs = v_abs
        self.v_r = v_r
        self.v_az = v_az
        self.v_el = v_el
        self._velocity_vecs_assigned = True


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import random
from SimInterface import SimInterface
from collections import namedtuple
import os
from tensorboardX import SummaryWriter
writer = SummaryWriter('./runs/exp1')
# Hyper Parameters
BATCH_SIZE = 2
LR = 0.01                   # learning rate
EPSILON = 0.9               # greedy policy
GAMMA = 0.9                 # reward discount
TARGET_REPLACE_ITER = 100
SAVEMOLDE_ITER=1000# target update frequency
SAVE_POSITION=0
MEMORY_CAPACITY = 10
env = SimInterface(agent_num=10, time=5, render=True)
N_ACTIONS = 14
N_STATES = 165
ENV_A_SHAPE = 0
ActionHash3 = [-1,0,1]
ActionHash2 = [0,1]
Transition = namedtuple('Transition', ['state', 'action', 'reward', 'next_state'])

def load_checkpoint(filename, model, optimizer):
    if os.path.isfile(filename):
        print("# loading checkpoint '{}'".format(filename))
        checkpoint = torch.load(filename)
        model.load_state_dict(checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        print("# loaded checkpoint '{}'".format(filename))
    else:
        print("# no checkpoint found at '{}'".format(filename))
    return model, optimizer

def save_checkpoint(filename, model, optimizer):
    state = {
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
    }
    torch.save(state, filename)

def log(ep, r, res_queue, global_nr_steps, steps):
    with ep.get_lock():
        ep.value += 1
    with ep_r.get_lock():
        if ep_r.value == 0.:
            ep_r.value = -1  # ep_r
            steps.value = steps
        else:
            ep_r.value = r.value * 0.99 + ep_r * 0.01
            global_nr_steps.value = global_nr_steps.value * 0.99 + steps * 0.01
    res_queue.put(ep_r)
    print(
        name,
        "Ep:", ep.value,
        "| Avg Ep_r: %.2f" % ep_r.value,
        "| Avg Steps: %d" % steps.value,
        "| Ep_r / Steps: %.2f" % (ep_r / steps),
    )

class ReplayMemory(object):
    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, t):
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = t
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)

class Net(nn.Module):
    def __init__(self, ):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(N_STATES, 900)
        self.fc1.weight.data.normal_(0, 0.1)   # initialization
        self.out = nn.Linear(900, N_ACTIONS)
        self.out.weight.data.normal_(0, 0.1)   # initialization

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        actions_value = self.out(x)
        return actions_value

class DQN(object):
    def __init__(self):
        self.eval_net, self.target_net = Net(), Net()
        self.learn_step_counter = 0                                     # for target updating
        self.memory_counter = 0                                         # for storing memory
        self.memory = np.zeros((MEMORY_CAPACITY, N_STATES * 2 + 2))     # initialize memory
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=LR)
        self.loss_func = nn.MSELoss()
        self.saveModle_counter=0
        self.evalfilename = './evalModel'
        self.targetfilename = './targetModel'

    def choose_action(self, x):
        #action shape = [0,0,0,0,0,0,0,0]
        action = []
        x = torch.unsqueeze(torch.FloatTensor(x), 0)
        # input only one sample
        if np.random.uniform() < EPSILON:   # greedy
            actions_value = self.eval_net.forward(x) # 20维
            actions_value=actions_value.data.numpy()
            a1 = ActionHash3[np.argmax(actions_value[:,0:3])]
            a2 = ActionHash3[np.argmax(actions_value[:, 3:6])]
            a3 = ActionHash3[np.argmax(actions_value[:,6:9])]
            a4 = ActionHash3[np.argmax(actions_value[:,9:12])]
            a5 = ActionHash2[np.argmax(actions_value[:,12:14])]
            action=[a1,a2,a3,a4,a5]
        else:   # random
            for i in range(4):
                in_action = np.random.randint(-1, 2)
                action.append(in_action)
            action.append(np.random.randint(0, 2))

        return action

    def store_transition(self, s, a, r, s_):
        transition = np.hstack((s, [a, r], s_))
        # replace the old memory with new memory
        index = self.memory_counter % MEMORY_CAPACITY
        self.memory[index, :] = transition
        self.memory_counter += 1

    def learn(self):

        # target parameter update
        if self.learn_step_counter % TARGET_REPLACE_ITER == 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())
        self.learn_step_counter += 1

        if self.saveModle_counter % SAVEMOLDE_ITER == 0:
            print("-----------------------MODLE---SAVED-------------------------------")
            self.evalfilename+=str(self.saveModle_counter/SAVEMOLDE_ITER)
            self.evalfilename +='.pth'
            self.targetfilename += str(self.saveModle_counter/SAVEMOLDE_ITER)
            self.targetfilename += '.pth'
            torch.save(self.eval_net.state_dict(), self.evalfilename)
            torch.save(self.target_net.state_dict(), self.targetfilename)
        self.saveModle_counter += 1

        # sample batch transitions
        # sample_index = np.random.choice(MEMORY_CAPACITY, BATCH_SIZE)
        # b_memory = self.memory[sample_index, :]
        b_memory= Memreplay.sample(BATCH_SIZE)
        b_s = torch.tensor([t.state for t in b_memory]).float()
        b_a = torch.tensor([t.action for t in b_memory]).long()
        b_r = torch.tensor([t.reward for t in b_memory]).float()
        b_s_ = torch.tensor([t.next_state for t in b_memory]).float()
        # b_s = torch.FloatTensor(b_memory[:, :N_STATES])
        # b_a = torch.LongTensor(b_memory[:, N_STATES:N_STATES+1].astype(int))
        # b_r = torch.FloatTensor(b_memory[:, N_STATES+1:N_STATES+2])
        # b_s_ = torch.FloatTensor(b_memory[:, -N_STATES:])
        # q_eval w.r.t the action in experience
        q_eval = self.eval_net(b_s).gather(1, b_a).sum(dim=1) # shape (batch, 1)
        q_next = self.target_net(b_s_).detach()     # detach from graph, don't backpropagate
        q_nextmax=q_next[:,0:3].max(1)[0]+q_next[:,3:6].max(1)[0]+\
                  q_next[:,6:9].max(1)[0]+q_next[:,9:12].max(1)[0]+q_next[:,12:14].max(1)[0]
                  # q_next[:,14:16].max(1)[0]+q_next[:,16:18].max(1)[0]+q_next[:,18:20].max(1)[0]
        q_target = b_r + GAMMA * q_nextmax  # shape (batch, 1)
        loss = self.loss_func(q_eval, q_target)
        self.optimizer.zero_grad()
        loss.backward()
        writer.add_scalar('loss', loss, self.learn_step_counter)
        self.optimizer.step()

Memreplay = ReplayMemory(MEMORY_CAPACITY)

def train():
    dqn = DQN()
    wintime = 0
    losetime = 0
    filename = './A3C_TP_trained_critic.pth'
    print('\nCollecting experience...')
    for i_episode in range(40000):
        s = env.reset()
        ep_steps = 0
        ep_r=-1.0
        while True:
            agent1s = np.concatenate((s, s[0:15]))
            agent2s = np.concatenate((s, s[15:30]))
            agent3s = np.concatenate((s, s[30:45]))
            agent4s = np.concatenate((s, s[45:60]))
            agent5s = np.concatenate((s, s[60:75]))
            random_action1 = []
            random_action2 = []
            random_action3 = []
            random_action4 = []
            random_action5 = []
            for i in range(8):
                in_action1 = np.random.randint(-1, 2)
                random_action1.append(in_action1)
                in_action2 = np.random.randint(-1, 2)
                random_action2.append(in_action2)
                in_action3 = np.random.randint(-1, 2)
                random_action3.append(in_action3)
                in_action4 = np.random.randint(-1, 2)
                random_action4.append(in_action4)
                in_action5 = np.random.randint(-1, 2)
                random_action5.append(in_action5)
            DNQ_action1 = dqn.choose_action(agent1s)
            DNQ_action2 = dqn.choose_action(agent2s)
            DNQ_action3 = dqn.choose_action(agent3s)
            DNQ_action4 = dqn.choose_action(agent4s)
            DNQ_action5 = dqn.choose_action(agent5s)
            DNQ_action1full = DNQ_action1 + [0, 1, 0]
            DNQ_action2full = DNQ_action2 + [0, 1, 0]
            DNQ_action3full = DNQ_action3 + [0, 1, 0]
            DNQ_action4full = DNQ_action4 + [0, 1, 0]
            DNQ_action5full = DNQ_action5 + [0, 1, 0]
            offset = np.array([1, 4, 7, 10, 12])
            a = np.array([DNQ_action1full, DNQ_action2full, DNQ_action3full, DNQ_action4full, DNQ_action5full,
                          random_action1, random_action2, random_action3, random_action4, random_action5])
            # take action
            s_, r, done, info = env.step(a)
            ep_steps+=1
            print(ep_steps)
            agent1nexts = np.concatenate((s_, s_[0:15]))
            agent2nexts = np.concatenate((s_, s_[15:30]))
            agent3nexts = np.concatenate((s_, s_[30:45]))
            agent4nexts = np.concatenate((s_, s_[45:60]))
            agent5nexts = np.concatenate((s_, s_[60:75]))
            transition = Transition(agent1s, DNQ_action1 + offset, r, agent1nexts)

            transition = Transition(agent2s, DNQ_action2 + offset, r, agent1nexts)
            transition = Transition(agent3s, DNQ_action3 + offset, r, agent1nexts)
            transition = Transition(agent4s, DNQ_action4 + offset, r, agent1nexts)
            transition = Transition(agent5s, DNQ_action5 + offset, r, agent1nexts)
            # dqn.store_transition(s, a, r, s_)
            Memreplay.push(transition)
            dqn.memory_counter += 5

            if dqn.memory_counter > MEMORY_CAPACITY:
                dqn.learn()
                if done:
                    if r==1:
                        str='GameWin'
                        wintime+=1
                    else:
                        str='Gamelose'
                        losetime+=1
                    ep_r = ep_r* 0.99 + r * 0.01
                    print("Episode:",i_episode,str,"Avg_Loss:",ep_r)

            if done:
                print("Episode Finish")
                break
            s = s_

def eval(world, init_gmodel=False):
    dqn=DQN()
    env = SimInterface(agent_num=10, time=180, render=True)

    do_print = True
    done = None
    reward = 0
    last_reward = [0, 0, 0, 0]

    while True:
        done, s = False, env.reset()
        agent1s = np.concatenate((s, s[0:15]))
        agent2s = np.concatenate((s, s[15:30]))
        agent3s = np.concatenate((s, s[30:45]))
        agent4s = np.concatenate((s, s[45:60]))
        agent5s = np.concatenate((s, s[60:75]))
        t = 0
        while True:
            random_action1 = []
            random_action2 = []
            random_action3 = []
            random_action4 = []
            random_action5 = []
            for i in range(8):
                in_action1 = np.random.randint(-1, 2)
                random_action1.append(in_action1)
                in_action2 = np.random.randint(-1, 2)
                random_action2.append(in_action2)
                in_action3 = np.random.randint(-1, 2)
                random_action3.append(in_action3)
                in_action4 = np.random.randint(-1, 2)
                random_action4.append(in_action4)
                in_action5 = np.random.randint(-1, 2)
                random_action5.append(in_action5)
            DNQ_action1 = dqn.choose_action(agent1s)
            DNQ_action2 = dqn.choose_action(agent2s)
            DNQ_action3 = dqn.choose_action(agent3s)
            DNQ_action4 = dqn.choose_action(agent4s)
            DNQ_action5 = dqn.choose_action(agent5s)
            DNQ_action1full = DNQ_action1 + [0, 1, 0]
            DNQ_action2full = DNQ_action2 + [0, 1, 0]
            DNQ_action3full = DNQ_action3 + [0, 1, 0]
            DNQ_action4full = DNQ_action4 + [0, 1, 0]
            DNQ_action5full = DNQ_action5 + [0, 1, 0]
            offset = np.array([1, 4, 7, 10, 12])
            a = np.array([DNQ_action1full, DNQ_action2full, DNQ_action3full, DNQ_action4full, DNQ_action5full,
                          random_action1, random_action2, random_action3, random_action4, random_action5])
            # take action
            s_, r, done, info = env.step(a)

if __name__ == '__main__':
    train()


import torch
import argparse
from scipy.io.wavfile import write
import numpy as np
from model.generator import ModifiedGenerator
from utils.hparams import HParam, load_hparam_str
from denoiser import Denoiser

MAX_WAV_VALUE = 32768.0

def main(args):
    checkpoint = torch.load(args.checkpoint_path)
    if args.config is not None:
        hp = HParam(args.config)
    else:
        hp = load_hparam_str(checkpoint['hp_str'])

    model = ModifiedGenerator(hp.audio.n_mel_channels, hp.model.n_residual_layers,
                        ratios=hp.model.generator_ratio, mult = hp.model.mult,
                        out_band = hp.model.out_channels).cuda()
    model.load_state_dict(checkpoint['model_g'])
    model.eval(inference=True)

    with torch.no_grad():
        mel = torch.from_numpy(np.load(args.input))
        if len(mel.shape) == 2:
            mel = mel.unsqueeze(0)
        mel = mel.cuda()
        zero = torch.full((1, 80, 10), -11.5129).to(mel.device)
        mel = torch.cat((mel, zero), dim=2)
        vocgan_trace = torch.jit.trace(model, mel)
        vocgan_trace.save("{}/vocgan_ex_female_en_{}_{}.pt".format(args.out, checkpoint['githash'], checkpoint['epoch']))
        # audio = model(mel)

        # audio = audio.squeeze(0)  # collapse all dimension except time axis
        # if args.d:
        #     denoiser = Denoiser(model).cuda()
        #     audio = denoiser(audio, 0.01)
        # audio = audio.squeeze()
        # audio = audio[:-(hp.audio.hop_length*10)]
        # audio = MAX_WAV_VALUE * audio
        # audio = audio.clamp(min=-MAX_WAV_VALUE, max=MAX_WAV_VALUE-1)
        # audio = audio.short()
        # audio = audio.cpu().detach().numpy()
        #
        # out_path = args.input.replace('.npy', '_reconstructed_epoch%04d.wav' % checkpoint['epoch'])
        # write(out_path, hp.audio.sampling_rate, audio)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str, default=None,
                        help="yaml file for config. will use hp_str from checkpoint if not given.")
    parser.add_argument('-p', '--checkpoint_path', type=str, required=True,
                        help="path of checkpoint pt file for evaluation")
    parser.add_argument('-i', '--input', type=str, required=True,
                        help="directory of mel-spectrograms to invert into raw audio. ")
    parser.add_argument('-o', '--out', type=str, required=True,
                        help="path of output pt file")
    args = parser.parse_args()

    main(args)


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

#
# SPDX-License-Identifier: GPL-3.0
#
# GNU Radio Python Flow Graph
# Title: Not titled yet
# GNU Radio version: 3.8.2.0

from gnuradio import blocks
import numpy
from gnuradio import gr
from gnuradio.filter import firdes
import sys
import signal
from argparse import ArgumentParser
from gnuradio.eng_arg import eng_float, intx
from gnuradio import eng_notation
import zynq

class loopback(gr.top_block):

    def __init__(self):
        gr.top_block.__init__(self, "Not titled yet")

        ##################################################
        # Variables
        ##################################################
        self.samp_rate = samp_rate = 32000

        ##################################################
        # Blocks
        ##################################################
        self.zynq_loopback_0 = zynq.loopback()
        self.blocks_file_sink_0 = blocks.file_sink(gr.sizeof_int*1, 'output.bin', False)
        self.blocks_file_sink_0.set_unbuffered(False)
        self.analog_random_source_x_0 = blocks.vector_source_i(list(map(int, numpy.random.randint(0, 2, 1*1024*1024))), False)

        ##################################################
        # Connections
        ##################################################
        self.connect((self.analog_random_source_x_0, 0), (self.zynq_loopback_0, 0))
        self.connect((self.zynq_loopback_0, 0), (self.blocks_file_sink_0, 0))

    def get_samp_rate(self):
        return self.samp_rate

    def set_samp_rate(self, samp_rate):
        self.samp_rate = samp_rate

def main(top_block_cls=loopback, options=None):
    tb = top_block_cls()

    def sig_handler(sig=None, frame=None):
        tb.stop()
        tb.wait()

        sys.exit(0)

    signal.signal(signal.SIGINT, sig_handler)
    signal.signal(signal.SIGTERM, sig_handler)

    tb.start()

    try:
        input('Press Enter to quit: ')
    except EOFError:
        pass
    tb.stop()
    tb.wait()

if __name__ == '__main__':
    main()


import cv2
import mxnet as mx
import numpy as np
from mxnet import ndarray as nd

IMG_WIDTH = 112
IMG_HIGHT = 112

class FaceAttributesPredictor:
    def __init__(self, prefix, epoch=0, ctxId=0):
        self.ctx = mx.gpu(ctxId)
        self.prefix = prefix
        self.epoch = epoch

        sym, arg_params, aux_params = mx.model.load_checkpoint(self.prefix, self.epoch)
        all_layers = sym.get_internals()
        # print(all_layers)
        # 提取中间的某一层作为输出，在里要取softmax层之前的全连接层作为输出
        sym = all_layers['fc1_output']
        self.model = mx.module.Module(symbol=sym, context=self.ctx, label_names=None)
        # self.model = mx.module.Module(symbol=sym, context=self.ctx, data_names=['data'], label_names=None)
        self.model.bind(for_training=False, data_shapes=[('data', (1, 3, IMG_WIDTH, IMG_HIGHT))])
        self.model.set_params(arg_params, aux_params)

    def predict(self, mats):
        assert len(mats) > 0
        # dataShape = (1, 3, IMG_WIDTH, IMG_HIGHT)
        # self.model.bind(for_training=False, data_shapes=[('data', dataShape)])

        result = []
        for i in range(len(mats)):
            if mats[i].shape != (IMG_WIDTH, IMG_HIGHT, 3):
                imgMat = cv2.resize(mats[i], (IMG_WIDTH, IMG_HIGHT))
            else:
                imgMat = mats[i]
            imgMat = imgMat.astype(np.float32)
            imTensor = np.zeros((1, 3, imgMat.shape[0], imgMat.shape[1]))
            for i in range(3):
                imTensor[0, i, :, :] = imgMat[:, :, 2 - i]  # bgr2rgb  mxnet rgb  opencv bgr and (h, w, c) to (c, h, w)

            data = nd.array(imTensor)
            db = mx.io.DataBatch(data=(data,), provide_data=[('data', data.shape)])

            self.model.forward(db, is_train=False)
            print('forward')
            # probs = self.mod.get_outputs()[0].asnumpy()
            probs = self.model.get_outputs()[0].asnumpy()
            print(probs)
            for prob in probs:
                pred = {}
                pred['Gender'] = np.argmax(prob[0:2])
                pred['Mask'] = np.argmax(prob[2:4])
                pred['Glass'] = np.argmax(prob[4:7])
                pred['MouthOpen'] = np.argmax(prob[7:10])
                pred['EyesClose'] = np.argmax(prob[10:13])
                result.append(pred)

        return result


"""
Add Wikipedia's category to feature.
"""
import argparse
import os
from collections import defaultdict

from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from sklearn.svm import LinearSVC
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, f1_score
from scipy.sparse import hstack

from entitypedia.classifier.preprocess import extract_abstract, tokenize, remove_tags, create_dictionary, \
    extract_categories
from entitypedia.corpora.datasets import DocumentClassifierDataset, load_disambig_ids, save_jsonl

def load_dataset(wiki_dir, seed_dir, categories):
    X, y = [], []
    cats = []
    dataset = DocumentClassifierDataset(wiki_dir, seed_dir)
    for j in dataset.create():
        try:
            text = extract_abstract(j['text'])
            text = remove_tags(text)
            X.append(text)
            y.append(j['label'])
            cats.append(' '.join(categories[j['id']]))
        except IndexError:
            pass

    return X, cats, y

def load_prediction_dataset(wiki_dir, ignored_ids, categories):
    X, ids = [], []
    cats = []
    titles = []
    dataset = DocumentClassifierDataset(wiki_dir)
    for j in dataset.create_prediction_set(ignored_ids):
        id = j['id']
        if id in ignored_ids:
            continue
        try:
            text = extract_abstract(j['text'])
        except IndexError:
            continue
        text = remove_tags(text)
        X.append(text)
        ids.append(id)
        cats.append(' '.join(categories[id]))
        titles.append(j['title'])

    return X, cats, ids, titles

def load_categories(file_path):
    with open(file_path) as f:
        lines = [line.strip().split('\t') for line in f]
    ids, cats = [], []
    for row in lines:
        ids.append(row[0])
        cats.append(row[1])
    cats = extract_categories(cats)
    res = defaultdict(list)
    for id, cat in zip(ids, cats):
        res[id].append(cat)

    return res

def main(args):
    print('Loading dataset...')
    categories = load_categories(args.category)
    X, cats, y = load_dataset(args.wiki_dir, args.seed_dir, categories)
    label_dict = create_dictionary([y])
    y = [label_dict.token2id[y_] for y_ in y]

    print('Vectorizing...')
    vectorizer = TfidfVectorizer(tokenizer=tokenize, analyzer='word')
    X = vectorizer.fit_transform(X)
    category_vectorizer = CountVectorizer()
    cats = category_vectorizer.fit_transform(cats)
    X = hstack([X, cats])

    # x_train, x_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

    print('Fitting...')
    clf = LinearSVC()
    clf.fit(X, y)
    # clf.fit(x_train, y_train)

    # print('Predicting...')
    # y_pred = clf.predict(x_test)

    # print('Evaluating...')
    # print(f1_score(y_test, y_pred, average='micro'))
    # print(f1_score(y_test, y_pred, average='macro'))
    # print(classification_report(y_test, y_pred, digits=3))

    print('Loading dataset for prediction...')
    disambig_ids = load_disambig_ids(args.disambig_file)
    X, cats, ids, titles = load_prediction_dataset(args.wiki_dir, disambig_ids, categories)

    print('Vectorizing...')
    X = vectorizer.transform(X)
    cats = category_vectorizer.transform(cats)
    X = hstack([X, cats])

    print('Predicting...')
    y_pred = clf.predict(X)

    print('Saving...')
    outputs = [{'id': int(id), 'ne_id': int(ne_id), 'title': title}
               for id, ne_id, title in zip(ids, y_pred, titles)]
    save_jsonl(outputs, args.save_file)
    label_dict.save(args.label_dic)

if __name__ == '__main__':
    DATA_DIR = os.path.join(os.path.dirname(__file__), '../../data/raw')
    SAVE_DIR = os.path.join(os.path.dirname(__file__), '../../data/interim')
    parser = argparse.ArgumentParser(description='Training a classifier')
    parser.add_argument('--wiki_dir', default=os.path.join(DATA_DIR, 'extracted'))
    parser.add_argument('--seed_dir', default=os.path.join(DATA_DIR, 'seeds'))
    parser.add_argument('--category', default=os.path.join(DATA_DIR, 'categories.tsv'))
    parser.add_argument('--disambig_file', default=os.path.join(DATA_DIR, 'disambig_id.csv'))
    parser.add_argument('--save_file', default=os.path.join(SAVE_DIR, 'pages.jsonl'))
    parser.add_argument('--label_dic', default=os.path.join(SAVE_DIR, 'labels.dic'))
    args = parser.parse_args()
    main(args)


import numpy as np
import torch
from gym import Wrapper
# from gym_maze.envs.maze_env import MazeEnvSample5x5

import wandb
from config import config
from embedding_model import compute_intrinsic_reward
from memory import LocalBuffer
from model import R2D2_agent57

# todo : MetaController 구현
class Maze(Wrapper):
    def step(self, action: int):
        obs, rew, done, info = super().step(["N", "E", "S", "W"][action])
        self.set.add((obs[0], obs[1]))
        if rew > 0:
            rew = 10
        return obs / 10, rew, done, info

    def reset(self):
        self.set = set()
        return super().reset()

def get_action(state, target_net, epsilon, env, hidden, beta):
    action, hidden = target_net.get_action(state, hidden, beta)

    if np.random.rand() <= epsilon:
        return env.action_space.sample(), hidden
    else:
        return action, hidden

import gym

class Actor:
    def __init__(self, actor_id, online_net, target_net, current_g_model, target_g_model, embedding_model, memory, meta_controller,
                 epsilon, lock):

        self.env = gym.make('CartPole-v1') #gym.Maze(MazeEnvSample5x5())
        # self.env = Maze(MazeEnvSample5x5())
        self.env.seed(config.random_seed)
        self.env.action_space.seed(config.random_seed)

        self.actor_id = actor_id
        self.online_net = online_net
        self.target_net = target_net
        self.current_g_model = current_g_model
        self.target_g_model = target_g_model
        self.embedding_model = embedding_model
        self.memory = memory
        self.mc = meta_controller
        self.epsilon = epsilon
        self.lock = lock

    def run(self):
        steps = 0
        local_buffer = LocalBuffer()

        for episode in range(30000):
            done = False
            sum_reward = 0
            sum_augmented_reward = 0
            state = self.env.reset()
            state = torch.Tensor(state).to(config.device)

            hidden = ((
                torch.Tensor().new_zeros(1, 1, config.hidden_size),
                torch.Tensor().new_zeros(1, 1, config.hidden_size),
            ), (
                torch.Tensor().new_zeros(1, 1, config.hidden_size),
                torch.Tensor().new_zeros(1, 1, config.hidden_size),
            )
            )

            episodic_memory = [self.embedding_model.embedding(state)]

            episode_steps = 0
            horizon = 100
            MA = 0.1

            while not done:
                steps += 1
                episode_steps += 1

                action, new_hidden = get_action(state, self.target_net, self.epsilon, self.env, hidden,
                                                beta=self.mc.beta)
                # print(f"self.actor_id={self.actor_id}, action={action}")

                next_state, env_reward, done, _ = self.env.step(action)
                next_state = torch.Tensor(next_state)

                augmented_reward = env_reward
                if config.enable_ngu:
                    next_state_emb = self.embedding_model.embedding(next_state)

                    c_out = self.current_g_model(next_state)
                    alpha = self.target_g_model.train_model(c_out, next_state)

                    intrinsic_reward, MA = compute_intrinsic_reward(episodic_memory, next_state_emb, alpha=alpha,
                                                                    episode_steps=episode_steps, MA=MA)

                    episodic_memory.append(next_state_emb)
                    beta = self.mc.beta
                    augmented_reward = env_reward + beta * intrinsic_reward

                mask = 0 if done else 1

                local_buffer.push(state.detach(), next_state.detach(), action, augmented_reward, mask, [hid_or_cell.detach() for hid_or_cell in hidden[0]],
                                  [hid_or_cell.detach() for hid_or_cell in hidden[1]], self.mc.gamma, self.mc.beta)
                hidden = new_hidden

                # todo :get_td_error 할 때 config.beta가 아니라 beta 가변적으로 받을 수 있도록
                if len(local_buffer.memory) == config.local_mini_batch:
                    batch, lengths = local_buffer.sample()
                    td_error = R2D2_agent57.get_td_error(self.online_net, self.target_net, batch, lengths)

                    self.lock.acquire()
                    print(f"td_error={td_error}")
                    self.memory.push(td_error.detach(), batch, lengths)
                    self.lock.release()

                sum_reward += env_reward
                state = next_state
                sum_augmented_reward += augmented_reward

            print(f"Ep {episode} / augmented sum reward : {sum_augmented_reward}, sum_reward : {sum_reward}")
            self.mc.reset_states(sum_reward)

            # if episode > 0 and episode % config.log_interval == 0:
            #     mean_reward = sum_reward / config.log_interval
            #     mean_augmented_reward = sum_augmented_reward / config.log_interval
            #     metrics = {
            #         "episode": episode,
            #         "mean_reward": mean_reward,
            #         "epsilon": self.epsilon,
            #         # "embedding_loss": embedding_loss,
            #         "loss": loss,
            #         "mean_augmented_reward": mean_augmented_reward,
            #         "steps": steps,
            #         "sum_obs_set": sum_obs_set / config.log_interval,
            #     }
            #     print(metrics)
            #     wandb.log(metrics)
            #
            #     sum_reward = 0
            #     sum_augmented_reward = 0
            #     sum_obs_set = 0