from __future__ import print_function, division

import os, os.path, sys, re, glob
import itertools
from copy import deepcopy
import logging
import json

from .config import on_rtd

if not on_rtd:
    import numpy as np
    import pandas as pd

    import numpy.random as rand
    from scipy.stats import gaussian_kde
    import scipy
    import emcee
    import corner

    try:
        import pymultinest
    except ImportError:
        logging.warning('PyMultiNest not imported.  MultiNest fits will not work.')

    import configobj
    from astropy.coordinates import SkyCoord

    try:
        basestring
    except NameError:
        basestring = str

from .utils import addmags
from .observation import ObservationTree, Observation, Source
from .priors import age_prior, distance_prior, AV_prior, q_prior
from .priors import salpeter_prior, feh_prior
from .isochrone import get_ichrone, Isochrone

def _parse_config_value(v):
    try:
        val = float(v)
    except:
        try:
            val = [float(x) for x in v]
        except:
            val = v
    #print('{} becomes {}, type={}'.format(v,val,type(val)))
    return val

class StarModel(object):
    """

    :param ic:
        :class:`Isochrone` object used to model star.

    :param obs: (optional)
        :class:`ObservationTree` object containing photometry information.
        If not provided, then one will be constructed from the provided
        keyword arguments (which must include at least one photometric
        bandpass).  This should only happen in the simplest case
        of a single star system---if multiple stars are detected
        in any of the observations being used, an :class:`ObservationTree`
        should be passed.  If `obs` is a string, then it is assumed
        to be a filename of an obs summary DataFrame.

    :param N:
        Number of model stars to assign to each "leaf node" of the
        :class:`ObservationTree`.  If you want to model a binary star,
        provide ``N=2``.

    :param **kwargs:
        Keyword arguments must be properties of given isochrone, e.g., logg,
        feh, Teff, and/or magnitudes.  The values represent measurements of
        the star, and must be in (value,error) format. All such keyword
        arguments will be held in ``self.properties``.  ``parallax`` is
        also a valid property, and should be provided in miliarcseconds,
        as is ``density`` [g/cc], and ``nu_max`` and ``delta_nu``
        (asteroseismic parameters in uHz.)
    """

    # These are allowable parameters that are not photometric bands
    _not_a_band = ('RA','dec','ra','Dec','maxAV','parallax',
                  'logg','Teff','feh','density', 'separation',
                 'PA','resolution','relative','N','index', 'id')

    _default_name = 'single'

    def __init__(self, ic, obs=None, N=1, index=0,
                 name='', use_emcee=False,
                 RA=None, dec=None, coords=None,
                 **kwargs):

        self.name = name if name else self._default_name

        if coords is None:
            if RA is not None and dec is not None:
                try:
                    coords = SkyCoord(RA, dec)
                except:
                    coords = SkyCoord(float(RA), float(dec), unit='deg')
        self.coords = coords
        self._ic = ic

        self.use_emcee = use_emcee

        # If obs is not provided, build it
        if obs is None:
            self._build_obs(**kwargs)
            self.obs.define_models(ic, N=N, index=index)
            self._add_properties(**kwargs)
        elif isinstance(obs, basestring):
            df = pd.read_csv(obs)
            obs = ObservationTree.from_df(df)
            obs.define_models(ic, N=N, index=index)
            self.obs = obs
            self._add_properties(**kwargs)
        else:
            self.obs = obs
            if len(self.obs.get_model_nodes())==0:
                self.obs.define_models(ic, N=N, index=index)
                self._add_properties(**kwargs)

        self._priors = {'mass':salpeter_prior,
                        'feh':feh_prior,
                        'q':q_prior,
                        'age':age_prior,
                        'distance':distance_prior,
                        'AV':AV_prior}

        self._bounds = {'mass':None,
                        'feh':None,
                        'age':None,
                        'q':q_prior.bounds,
                        'distance':distance_prior.bounds,
                        'AV':AV_prior.bounds}

        if 'maxAV' in kwargs:
            self.set_bounds(AV=(0, kwargs['maxAV']))

        if 'max_distance' in kwargs:
            self.set_bounds(distance=(0, kwargs['max_distance']))

        self._directory = '.'
        self._samples = None

    @property
    def directory(self):
        return self._directory

    @property
    def ic(self):
        if type(self._ic)==type:
            self._ic = self._ic()
        return self._ic

    @classmethod
    def _parse_band(cls, kw):
        """Returns photometric band from inifile keyword
        """
        m = re.search('([a-zA-Z0-9]+)(_\d+)?', kw)
        if m:
            if m.group(1) in cls._not_a_band:
                return None
            else:
                return m.group(1)

    @classmethod
    def get_bands(cls, inifile):

        bands = []
        c = configobj.ConfigObj(inifile)
        for kw,v in c.items():
            if type(v) is configobj.Section:
                for kw in v:
                    b = cls._parse_band(kw)
                    if b is not None:
                        bands.append(b)
            else:
                b = cls._parse_band(kw)
                if b is not None:
                    bands.append(b)

        return list(set(bands))

    @classmethod
    def from_ini(cls, ic, folder='.', ini_file='star.ini', **kwargs):
        """
        Initialize a StarModel from a .ini file

        The "classic" format (version <= 0.9) should still work for a single star,
        where all properties are just listed in the file; e.g.,

            J = 10, 0.05
            H = 9.5, 0.05
            K = 9.0, 0.05
            Teff = 5000, 150

        If there are multiple stars observed, you can either define them in
        the ini file, or use the `obsfile` keyword, pointing to a file with
        the summarized photometric observations.  In this case, spectroscopic/parallax
        info should still be included in the .ini file; e.g.,

            obsfile = obs.csv
            Teff = 5000, 150

        The obsfile should be a comma-separated table with the following columns:
        `[name, band, resolution, mag, e_mag, separation, pa, relative]`.

          * `name` is the name of instrument
          * `band` is the photometric bandpass
          * `resolution` is the approximate spatial resolution of instrument
          * `mag`, `e_mag` describe magnitude of source (absolute or relative)
          * `separation`, `pa` describe position of source
          * `relative`: single-bit flag; if 1 then magnitudes taken with this
            instrument are assumed to be relative rather than absolute.

        If an obsfile is not provided, you can also define all the same information
        in the ini file, following these rules:

          * Every instrument/survey gets its own [section].  Sections are only
        created for different photometric observations.

          * if photometry relates to *all* stars in aperture,
            there is no extra info in the section, just the photometry.  In this case, it is
            also assumed that the photometry is absolute. (`relative=False`)

          * If 'resolution' is an attribute under a particular survey section (and
        'relative' is not explicitly stated), then the survey is assumed to have relative
        photometry, and to be listing
        information about companion stars.  In this case, there must be "separation"
        and "PA" included for each companion.  If there is more than one companion star,
        they must be identifed by tag, e.g., separation_1, PA_1, Ks_1, J_1, etc.  The
        tag can be anything alphanumeric, but it must be consistent within a particular
        section (instrument).  If there
        is no tag, there is assumed to be only one companion detected.

          * If there are no sections, then bands will be interpreted at face value
        and will all be assumed to apply to all stars modeled.

          * Default is to model each star in the highest-resolution observation as a
        single star, at the same distance/age/feh/AV.

        The `N` and `index`
        parameters may also be provided, to specify the relations between the
        model stars.  If these are not provided, then `N` will default to `1`
        (one model star per star observed in highest-resolution observation)
        and `index` will default to all `0` (all stars physically associated).

        """

        if not os.path.isabs(ini_file):
            ini_file = os.path.join(folder,ini_file)

        bands = cls.get_bands(ini_file)

        if not isinstance(ic, Isochrone):
            ic = get_ichrone(ic, bands)

        logging.debug('Initializing StarModel from {}'.format(ini_file))

        c = configobj.ConfigObj(ini_file)

        RA = c.get('RA')
        dec = c.get('dec')
        maxAV = c.get('maxAV')

        if len(c.sections) == 0:
            for k in c:
                kwargs[k] = _parse_config_value(c[k])
            obs = None
        else:

            columns = ['name', 'band', 'resolution', 'relative', 'separation', 'pa', 'mag', 'e_mag']
            df = pd.DataFrame(columns=columns)
            i = 0
            for k in c:
                if type(c[k]) != configobj.Section:
                    kwargs[k] = _parse_config_value(c[k])
                else:
                    instrument = k

                    # Set values of 'resolution' and 'relative'
                    if 'resolution' in c[k]:
                        resolution = float(c[k]['resolution'])
                        relative = True
                    else:
                        resolution = 4.0 #default
                        relative = False

                    # Overwrite value of 'relative' if it is explicitly set
                    if 'relative' in c[k]:
                        relative = c[k]['relative']=='True'

                    # Check if there are multiple stars (defined by whether
                    # any separations are listed).
                    # While we're at it, keep track of tags if they exist,
                    #  and pull out the names of the bands.
                    multiple = False
                    tags = []
                    bands = []
                    for label in c[k]:
                        m = re.search('separation(_\w+)?', label)
                        if m:
                            multiple = True
                            if m.group(1) is not None:
                                if m.group(1) not in tags:
                                    tags.append(m.group(1))
                        elif re.search('PA', label) or re.search('id', label) or \
                                label in ['resolution', 'relative']:
                            continue
                        else:
                            # At this point, this should be a photometric band
                            m = re.search('([a-zA-Z0-9]+)(_\w+)?', label)
                            b = m.group(1)
                            if b not in bands:
                                bands.append(b)

                    # If a blank tags needs to be created, do so
                    if len(bands) > 0 and (len(tags)==0 or bands[0] in c[k]):
                        tags.append('')

                    # For each band and each star, create a row
                    for b in bands:
                        for tag in tags:
                            if '{}{}'.format(b, tag) not in c[k]:
                                continue
                            row = {}
                            row['name'] = instrument
                            row['band'] = b
                            row['resolution'] = resolution
                            row['relative'] = relative
                            if 'separation{}'.format(tag) in c[k]:
                                row['separation'] = c[k]['separation{}'.format(tag)]
                                row['pa'] = c[k]['PA{}'.format(tag)]
                            else:
                                row['separation'] = 0.
                                row['pa'] = 0.
                            mag, e_mag = c[k]['{}{}'.format(b,tag)]
                            row['mag'] = float(mag)
                            row['e_mag'] = float(e_mag)
                            if not np.isnan(row['mag']) and not np.isnan(row['e_mag']):
                                df = df.append(pd.DataFrame(row, index=[i]))
                                i += 1

                        # put the reference star in w/ mag=0
                        if relative:
                            row = {}
                            row['name'] = instrument
                            row['band'] = b
                            row['resolution'] = resolution
                            row['relative'] = relative
                            row['separation'] = 0.
                            row['pa'] = 0.
                            row['mag'] = 0.
                            row['e_mag'] = 0.01
                            df = df.append(pd.DataFrame(row, index=[i]))
                            i += 1

            obs = ObservationTree.from_df(df)

        if 'obsfile' in c:
            obs = c['obsfile']

        logging.debug('Obs is {}'.format(obs))

        if 'name' not in kwargs:
            kwargs['name'] = os.path.basename(folder)
        new = StarModel(ic, obs=obs, **kwargs)
        new._directory = os.path.abspath(folder)

        return new

    def print_ascii(self):
        """Prints an ascii representation of the observation tree structure.
        """
        return self.obs.print_ascii()

    def bounds(self, prop):
        if self._bounds[prop] is not None:
            return self._bounds[prop]
        elif prop=='mass':
            lo, hi = (self.ic.minmass, self.ic.maxmass)
            self._bounds['mass'] = (lo, hi)
            self._priors['mass'].bounds = (lo, hi)
        elif prop=='feh':
            lo, hi = (self.ic.minfeh, self.ic.maxfeh)
            self._bounds['feh'] = (lo, hi)
            self._priors['feh'].bounds = (lo, hi)
        elif prop=='age':
            lo, hi = (self.ic.minage, self.ic.maxage)
            self._bounds['age'] = (lo, hi)
            self._priors['age'].bounds = (lo, hi)
            self._bounds['age'] = (self.ic.minage,
                                   self.ic.maxage)
        else:
            raise ValueError('Unknown property {}'.format(prop))
        return self._bounds[prop]

    def set_bounds(self, **kwargs):
        for k,v in kwargs.items():
            if len(v) != 2:
                raise ValueError('Must provide (min, max)')
            self._bounds[k] = v
            self._priors[k].bounds = v

    def _build_obs(self, **kwargs):
        """
        Builds ObservationTree out of keyword arguments

        Ignores anything that is not a photometric bandpass.
        This should not be used if there are multiple stars observed.

        Creates self.obs
        """
        logging.debug('Building ObservationTree...')
        tree = ObservationTree()
        for k,v in kwargs.items():
            if k in self.ic.bands:
                if np.size(v) != 2:
                    logging.warning('{}={} ignored.'.format(k,v))
                    # continue
                    v = [v, np.nan]
                o = Observation('', k, 99) #bogus resolution=99
                s = Source(v[0], v[1])
                o.add_source(s)
                logging.debug('Adding {} ({})'.format(s,o))
                tree.add_observation(o)

        self.obs = tree

    def _add_properties(self, **kwargs):
        """
        Adds non-photometry properties to ObservationTree
        """
        for k,v in kwargs.items():
            if k=='parallax':
                self.obs.add_parallax(v)
            elif k in ['Teff','logg','feh']:
                par = {k:v}
                self.obs.add_spectroscopy(**par)
            elif re.search('_', k):
                m = re.search('^(\w+)_(\w+)$', k)
                prop = m.group(1)
                tag = m.group(2)
                self.obs.add_spectroscopy(**{prop:v, 'label':'0_{}'.format(tag)})

    @property
    def param_description(self):
        return self.obs.param_description

    @property
    def param_names(self):
        return self.param_description

    @property
    def mags(self):
        return {n.band : n.value[0] for n in self.obs.get_obs_nodes()}

    def lnpost(self, p, **kwargs):
        lnpr = self.lnprior(p)
        if not np.isfinite(lnpr):
            return lnpr
        return lnpr + self.lnlike(p, **kwargs)

    def lnlike(self, p, **kwargs):
        lnl = self.obs.lnlike(p, **kwargs)
        return lnl

    def lnprior(self, p):
        N = self.obs.Nstars
        i = 0
        lnp = 0
        for s in self.obs.systems:
            age, feh, dist, AV = p[i+N[s]:i+N[s]+4]
            for prop, val in zip(['age','feh','distance','AV'],
                                 [age, feh, dist, AV]):
                lo,hi = self.bounds(prop)
                if val < lo or val > hi:
                    return -np.inf
                lnp += np.log(self.prior(prop, val))
                if not np.isfinite(lnp):
                    logging.debug('lnp=-inf for {}={} (system {})'.format(prop,val,s))
                    return -np.inf

            # Note: this is just assuming proper order.
            #  Is this OK?  Should keep eye out for bugs here.

            masses = p[i:i+N[s]]

            # Mass prior for primary
            lnp += np.log(self.prior('mass', masses[0]))
            if not np.isfinite(lnp):
                logging.debug('lnp=-inf for mass={} (system {})'.format(masses[0],s))

            # Priors for mass ratios
            for j in range(N[s]-1):
                q = masses[j+1]/masses[0]
                qmin, qmax = self.bounds('q')

                ## The following would enforce MA > MB > MC, but seems to make things very slow:
                #if j+1 > 1:
                #    qmax = masses[j] / masses[0]

                lnp += np.log(self.prior('q', q))
                if not np.isfinite(lnp):
                    logging.debug('lnp=-inf for q={} (system {})'.format(q,s))
                    return -np.inf

            i += N[s] + 4

        return lnp

    def prior(self, prop, val, **kwargs):
        return self._priors[prop](val, **kwargs)

    @property
    def n_params(self):
        tot = 0
        for _,n in self.obs.Nstars.items():
            tot += 4+n
        return tot

    def mnest_prior(self, cube, ndim, nparams):
        i = 0
        for _,n in self.obs.Nstars.items():
            minmass, maxmass = self.bounds('mass')
            for j in range(n):
                cube[i+j] = (maxmass - minmass)*cube[i+j] + minmass

            for j, par in enumerate(['age','feh','distance','AV']):
                lo, hi = self.bounds(par)
                cube[i+n+j] = (hi - lo)*cube[i+n+j] + lo
            i += 4 + n

    def mnest_loglike(self, cube, ndim, nparams):
        """loglikelihood function for multinest
        """
        return self.lnpost(cube)

    @property
    def labelstring(self):
        return '--'.join(['-'.join([n.label for n in l.children]) for l in self.obs.get_obs_leaves()])

    def fit(self, **kwargs):
        if self.use_emcee:
            return self.fit_mcmc(**kwargs)
        else:
            return self.fit_multinest(**kwargs)

    @property
    def mnest_basename(self):
        """Full path to basename
        """
        if not hasattr(self, '_mnest_basename'):
            s = self.labelstring
            if s=='0_0':
                s = 'single'
            elif s=='0_0-0_1':
                s = 'binary'
            elif s=='0_0-0_1-0_2':
                s = 'triple'

            s = '{}-{}'.format(self.ic.name, s)
            self._mnest_basename = os.path.join('chains', s+'-')

        if os.path.isabs(self._mnest_basename):
            return self._mnest_basename
        else:
            return os.path.join(self.directory, self._mnest_basename)

    @mnest_basename.setter
    def mnest_basename(self, basename):
        if os.path.isabs(basename):
            self._mnest_basename = basename
        else:
            self._mnest_basename = os.path.join('chains', basename)

    def lnpost_polychord(self, theta):
        phi = [0.0] #nDerived
        return self.lnpost(theta), phi

    def fit_polychord(self, basename, verbose=False, **kwargs):
        from .config import POLYCHORD
        sys.path.append(POLYCHORD)
        import PyPolyChord.PyPolyChord as PolyChord

        return PolyChord.run_nested_sampling(self.lnpost_polychord,
                        self.n_params, 0, file_root=basename, **kwargs)

    def fit_multinest(self, n_live_points=1000, basename=None,
                      verbose=True, refit=False, overwrite=False,
                      test=False,
                      **kwargs):
        """
        Fits model using MultiNest, via pymultinest.

        :param n_live_points:
            Number of live points to use for MultiNest fit.

        :param basename:
            Where the MulitNest-generated files will live.
            By default this will be in a folder named `chains`
            in the current working directory.  Calling this
            will define a `_mnest_basename` attribute for
            this object.

        :param verbose:
            Whether you want MultiNest to talk to you.

        :param refit, overwrite:
            Set either of these to true if you want to
            delete the MultiNest files associated with the
            given basename and start over.

        :param **kwargs:
            Additional keyword arguments will be passed to
            :func:`pymultinest.run`.

        """

        if basename is not None: #Should this even be allowed?
            self.mnest_basename = basename

        basename = self.mnest_basename
        if verbose:
            logging.info('MultiNest basename: {}'.format(basename))

        folder = os.path.abspath(os.path.dirname(basename))
        if not os.path.exists(folder):
            os.makedirs(folder)

        #If previous fit exists, see if it's using the same
        # observed properties
        prop_nomatch = False
        propfile = '{}properties.json'.format(basename)

        """
        if os.path.exists(propfile):
            with open(propfile) as f:
                props = json.load(f)
            if set(props.keys()) != set(self.properties.keys()):
                prop_nomatch = True
            else:
                for k,v in props.items():
                    if np.size(v)==2:
                        if not self.properties[k][0] == v[0] and \
                                self.properties[k][1] == v[1]:
                            props_nomatch = True
                    else:
                        if not self.properties[k] == v:
                            props_nomatch = True

        if prop_nomatch and not overwrite:
            raise ValueError('Properties not same as saved chains ' +
                            '(basename {}*). '.format(basename) +
                            'Use overwrite=True to fit.')
        """

        if refit or overwrite:
            files = glob.glob('{}*'.format(basename))
            [os.remove(f) for f in files]

        short_basename = self._mnest_basename

        mnest_kwargs = dict(n_live_points=n_live_points, outputfiles_basename=short_basename,
                        verbose=verbose)

        for k,v in kwargs.items():
            mnest_kwargs[k] = v

        if test:
            print('pymultinest.run() with the following kwargs: {}'.format(mnest_kwargs))
        else:
            wd = os.getcwd()
            os.chdir(os.path.join(folder, '..'))
            pymultinest.run(self.mnest_loglike, self.mnest_prior, self.n_params,
                            **mnest_kwargs)
            os.chdir(wd)
            #with open(propfile, 'w') as f:
            #    json.dump(self.properties, f, indent=2)

            self._make_samples()

    @property
    def mnest_analyzer(self):
        """
        PyMultiNest Analyzer object associated with fit.

        See PyMultiNest documentation for more.
        """
        return pymultinest.Analyzer(self.n_params, self.mnest_basename)

    @property
    def evidence(self):
        """
        Log(evidence) from multinest fit
        """
        s = self.mnest_analyzer.get_stats()
        return (s['global evidence'],s['global evidence error'])

    def maxlike(self, p0, **kwargs):
        """ Finds (local) optimum in parameter space.
        """
        def fn(p):
            return -self.lnpost(p)

        if 'method' not in kwargs:
            kwargs['method'] = 'Nelder-Mead'

        p0 = [0.8, 9.5, 0.0, 200, 0.2]
        fit = scipy.optimize.minimize(fn, p0, **kwargs)
        return fit

    def sample_from_prior(self, n):
        return self.emcee_p0(n)

    def emcee_p0(self, nwalkers):

        def sample_row(nstars, n=nwalkers):
            p = []
            m0 = self._priors['mass'].sample(n)
            age0 = self._priors['age'].sample(n)
            feh0 = self._priors['feh'].sample(n)
            d0 = self._priors['distance'].sample(n)
            AV0 = self._priors['AV'].sample(n)

            for i in range(nstars):
                p += [m0 * 0.95**i]
            p += [age0, feh0, d0, AV0]
            return p

        p0 = []
        for _,n in self.obs.Nstars.items():
            p0 += sample_row(n)

        p0 = np.array(p0).T

        nbad = 1

        while True:
            ibad = []
            for i, p in enumerate(p0):
                if not np.isfinite(self.lnpost(p)):
                    ibad.append(i)

            nbad = len(ibad)
            if nbad == 0:
                break

            pnew = []
            for _, n in self.obs.Nstars.items():
                pnew += sample_row(n, n=nbad)

            pnew = np.array(pnew).T

            p0[ibad, :] = pnew

        return p0

    def fit_mcmc(self,nwalkers=300,nburn=200,niter=100,
                 p0=None,initial_burn=None,
                 ninitial=50, loglike_kwargs=None,
                 **kwargs):
        """Fits stellar model using MCMC.

        :param nwalkers: (optional)
            Number of walkers to pass to :class:`emcee.EnsembleSampler`.
            Default is 200.

        :param nburn: (optional)
            Number of iterations for "burn-in."  Default is 100.

        :param niter: (optional)
            Number of for-keeps iterations for MCMC chain.
            Default is 200.

        :param p0: (optional)
            Initial parameters for emcee.  If not provided, then chains
            will behave according to whether inital_burn is set.

        :param initial_burn: (optional)
            If `True`, then initialize walkers first with a random initialization,
            then cull the walkers, keeping only those with > 15% acceptance
            rate, then reinitialize sampling.  If `False`, then just do
            normal burn-in.  Default is `None`, which will be set to `True` if
            fitting for distance (i.e., if there are apparent magnitudes as
            properties of the model), and `False` if not.

        :param ninitial: (optional)
            Number of iterations to test walkers for acceptance rate before
            re-initializing.

        :param loglike_args:
            Any arguments to pass to :func:`StarModel.loglike`, such
            as what priors to use.

        :param **kwargs:
            Additional keyword arguments passed to :class:`emcee.EnsembleSampler`
            constructor.

        :return:
            :class:`emcee.EnsembleSampler` object.

        """

        #clear any saved _samples
        if self._samples is not None:
            self._samples = None

        npars = self.n_params

        if p0 is None:
            p0 = self.emcee_p0(nwalkers)
            if initial_burn:
                sampler = emcee.EnsembleSampler(nwalkers,npars,self.lnpost,
                                                **kwargs)
                #ninitial = 300 #should this be parameter?
                pos, prob, state = sampler.run_mcmc(p0, ninitial)

                # Choose walker with highest final lnprob to seed new one
                i,j = np.unravel_index(sampler.lnprobability.argmax(),
                                        sampler.shape)
                p0_best = sampler.chain[i,j,:]
                print("After initial burn, p0={}".format(p0_best))
                p0 = p0_best * (1 + rand.normal(size=p0.shape)*0.001)
                print(p0)
        else:
            p0 = np.array(p0)
            p0 = rand.normal(size=(nwalkers,npars))*0.01 + p0.T[None,:]

        sampler = emcee.EnsembleSampler(nwalkers,npars,self.lnpost)
        pos, prob, state = sampler.run_mcmc(p0, nburn)
        sampler.reset()
        sampler.run_mcmc(pos, niter, rstate0=state)

        self._sampler = sampler
        return sampler

    @property
    def sampler(self):
        """
        Sampler object from MCMC run.
        """
        if hasattr(self,'_sampler'):
            return self._sampler
        else:
            raise AttributeError('MCMC must be run to access sampler')

    def _make_samples(self):

        if not self.use_emcee:
            filename = '{}post_equal_weights.dat'.format(self.mnest_basename)
            try:
                chain = np.loadtxt(filename)
                try:
                    lnprob = chain[:,-1]
                    chain = chain[:,:-1]
                except IndexError:
                    lnprob = np.array([chain[-1]])
                    chain = np.array([chain[:-1]])
            except:
                logging.error('Error loading chains from {}'.format(filename))
                raise
        else:
            #select out only walkers with > 0.15 acceptance fraction
            ok = self.sampler.acceptance_fraction > 0.15

            chain = self.sampler.chain[ok,:,:]
            chain = chain.reshape((chain.shape[0]*chain.shape[1],
                                        chain.shape[2]))

            lnprob = self.sampler.lnprobability[ok, :].ravel()

        df = pd.DataFrame()

        i=0
        for s,n in self.obs.Nstars.items():
            age = chain[:,i+n]
            feh = chain[:,i+n+1]
            distance = chain[:,i+n+2]
            AV = chain[:,i+n+3]
            for j in range(n):
                mass = chain[:,i+j]
                d = self.ic(mass, age, feh,
                             distance=distance, AV=AV)
                for c in d.columns:
                    df[c+'_{}_{}'.format(s,j)] = d[c]
            df['age_{}'.format(s)] = age
            df['feh_{}'.format(s)] = feh
            df['distance_{}'.format(s)] = distance
            df['AV_{}'.format(s)] = AV

            i += 4 + n

        for b in self.ic.bands:
            tot = np.inf
            for s,n in self.obs.Nstars.items():
                for j in range(n):
                    tot = addmags(tot,df[b + '_mag_{}_{}'.format(s,j)])
            df[b + '_mag'] = tot

        df['lnprob'] = lnprob

        self._samples = df.copy()

    @property
    def samples(self):
        """Dataframe with samples drawn from isochrone according to posterior

        Columns include both the sampling parameters from the MCMC
        fit (mass, age, Fe/H, [distance, A_V]), and also evaluation
        of the :class:`Isochrone` at each of these sample points---this
        is how chains of physical/observable parameters get produced.

        """
        if not hasattr(self,'sampler') and self._samples is None:
            raise AttributeError('Must run MCMC (or load from file) '+
                                 'before accessing samples')

        if self._samples is not None:
            df = self._samples
        else:
            self._make_samples()
            df = self._samples

        return df

    def random_samples(self, n):
        """
        Returns a random sampling of given size from the existing samples.

        :param n:
            Number of samples

        :return:
            :class:`pandas.DataFrame` of length ``n`` with random samples.
        """
        samples = self.samples
        inds = rand.randint(len(samples),size=int(n))

        newsamples = samples.iloc[inds]
        newsamples.reset_index(inplace=True)
        return newsamples

    def triangle(self, *args, **kwargs):
        return self.corner(*args, **kwargs)

    def corner(self, params, query=None, **kwargs):
        df = self.samples
        if query is not None:
            df = df.query(query)

        priors = []
        for p in params:
            if re.match('mass', p):
                priors.append(lambda x: self.prior('mass', x, bounds=self.bounds('mass')))
            elif re.match('age', p):
                priors.append(lambda x: self.prior('age', x, bounds=self.bounds('age')))
            elif re.match('feh', p):
                priors.append(lambda x: self.prior('feh', x, bounds=self.bounds('feh')))
            elif re.match('distance', p):
                priors.append(lambda x: self.prior('distance', x, bounds=self.bounds('distance')))
            elif re.match('AV', p):
                priors.append(lambda x: self.prior('AV', x, bounds=self.bounds('AV')))
            else:
                priors.append(None)

        try:
            fig = corner.corner(df[params], labels=params, priors=priors, **kwargs)
        except:
            logging.warning("Use Tim's version of corner to plot priors.")
            fig = corner.corner(df[params], labels=params, **kwargs)
        fig.suptitle(self.name, fontsize=22)
        return fig

    def triangle_physical(self, *args, **kwargs):
        return self.corner_physical(*args, **kwargs)

    def corner_plots(self, basename, **kwargs):
        fig1, fig2 = self.corner_physical(**kwargs), self.corner_observed(**kwargs)
        fig1.savefig(basename + '_physical.png')
        fig2.savefig(basename + '_observed.png')
        return fig1, fig2

    def triangle_plots(self, *args, **kwargs):
        return self.corner_plots(*args, **kwargs)

    def corner_physical(self, props=['mass','radius','feh','age','distance','AV'], **kwargs):
        collective_props = ['feh','age','distance','AV']
        indiv_props = [p for p in props if p not in collective_props]
        sys_props = [p for p in props if p in collective_props]

        props = ['{}_{}'.format(p,l) for p in indiv_props for l in self.obs.leaf_labels]
        props += ['{}_{}'.format(p,s) for p in sys_props for s in self.obs.systems]

        if 'range' not in kwargs:
            rng = [0.995 for p in props]

        return self.corner(props, range=rng, **kwargs)

    def mag_plot(self, *args, **kwargs):
        pass

    def corner_observed(self, **kwargs):
        """Makes corner plot for each observed node magnitude
        """
        tot_mags = []
        names = []
        truths = []
        rng = []
        for n in self.obs.get_obs_nodes():
            labels = [l.label for l in n.get_model_nodes()]
            band = n.band
            mags = [self.samples['{}_mag_{}'.format(band, l)] for l in labels]
            tot_mag = addmags(*mags)

            if n.relative:
                name = '{} $\Delta${}'.format(n.instrument, n.band)
                ref = n.reference
                if ref is None:
                    continue
                ref_labels = [l.label for l in ref.get_model_nodes()]
                ref_mags = [self.samples['{}_mag_{}'.format(band, l)] for l in ref_labels]
                tot_ref_mag = addmags(*ref_mags)
                tot_mags.append(tot_mag - tot_ref_mag)
                truths.append(n.value[0] - ref.value[0])
            else:
                name = '{} {}'.format(n.instrument, n.band)
                tot_mags.append(tot_mag)
                truths.append(n.value[0])

            names.append(name)
            rng.append((min(truths[-1], np.percentile(tot_mags[-1],0.5)),
                        max(truths[-1], np.percentile(tot_mags[-1],99.5))))
        tot_mags = np.array(tot_mags).T

        return corner.corner(tot_mags, labels=names, truths=truths, range=rng, **kwargs)

    def save_hdf(self, filename, path='', overwrite=False, append=False):
        """Saves object data to HDF file (only works if MCMC is run)

        Samples are saved to /samples location under given path,
        :class:`ObservationTree` is saved to /obs location under given path.

        :param filename:
            Name of file to save to.  Should be .h5 file.

        :param path: (optional)
            Path within HDF file structure to save to.

        :param overwrite: (optional)
            If ``True``, delete any existing file by the same name
            before writing.

        :param append: (optional)
            If ``True``, then if a file exists, then just the path
            within the file will be updated.
        """
        if os.path.exists(filename):
            with pd.HDFStore(filename) as store:
                if path in store:
                    if overwrite:
                        os.remove(filename)
                    elif not append:
                        raise IOError('{} in {} exists.  Set either overwrite or append option.'.format(path,filename))

        if self.samples is not None:
            self.samples.to_hdf(filename, path+'/samples')
        else:
            pd.DataFrame().to_hdf(filename, path+'/samples')

        self.obs.save_hdf(filename, path+'/obs', append=True)

        with pd.HDFStore(filename) as store:
            # store = pd.HDFStore(filename)
            attrs = store.get_storer('{}/samples'.format(path)).attrs

            attrs.ic_type = type(self.ic)
            attrs.ic_bands = list(self.ic.bands)
            attrs.use_emcee = self.use_emcee
            if hasattr(self, '_mnest_basename'):
                attrs._mnest_basename = self._mnest_basename

            attrs._bounds = self._bounds
            attrs._priors = self._priors

            attrs.name = self.name
            store.close()

    @classmethod
    def load_hdf(cls, filename, path='', name=None):
        """
        A class method to load a saved StarModel from an HDF5 file.

        File must have been created by a call to :func:`StarModel.save_hdf`.

        :param filename:
            H5 file to load.

        :param path: (optional)
            Path within HDF file.

        :return:
            :class:`StarModel` object.
        """
        if not os.path.exists(filename):
            raise IOError('{} does not exist.'.format(filename))
        store = pd.HDFStore(filename)
        try:
            samples = store[path+'/samples']
            attrs = store.get_storer(path+'/samples').attrs
        except:
            store.close()
            raise

        try:
            ic = attrs.ic_type(attrs.ic_bands)
        except AttributeError:
            ic = attrs.ic_type

        use_emcee = attrs.use_emcee
        mnest = True
        try:
            basename = attrs._mnest_basename
        except AttributeError:
            mnest = False
        bounds = attrs._bounds
        priors = attrs._priors

        if name is None:
            try:
                name = attrs.name
            except:
                name = ''

        store.close()

        obs = ObservationTree.load_hdf(filename, path+'/obs', ic=ic)

        mod = cls(ic, obs=obs,
                  use_emcee=use_emcee, name=name)
        mod._samples = samples
        if mnest:
            mod._mnest_basename = basename
        mod._directory = os.path.dirname(filename)
        return mod

class BinaryStarModel(StarModel):
    _default_name = 'binary'
    def __init__(self, *args, **kwargs):
        kwargs['N'] = 2
        super(BinaryStarModel, self).__init__(*args, **kwargs)

    @classmethod
    def from_ini(cls, *args, **kwargs):
        kwargs['N'] = 2
        return super(BinaryStarModel, cls).from_ini(*args, **kwargs)

class TripleStarModel(StarModel):
    _default_name = 'triple'
    def __init__(self, *args, **kwargs):
        kwargs['N'] = 3
        super(TripleStarModel, self).__init__(*args, **kwargs)

    @classmethod
    def from_ini(cls, *args, **kwargs):
        kwargs['N'] = 3
        return super(TripleStarModel, cls).from_ini(*args, **kwargs)

class StarModelGroup(object):
    """A collection of StarModel objects with different model node specifications

    Pass a single StarModel, and model nodes will be cleared and replaced with
    different variants.
    """
    def __init__(self, base_model, max_multiples=1, max_stars=2):

        self.base_model = deepcopy(base_model)
        self.base_model.obs.clear_models()
        self.max_multiples = max_multiples
        self.max_stars = max_stars

        self.models = []
        for N, index in self.model_options:
            mod = deepcopy(self.base_model)
            mod.obs.define_models(self.ic, N=N, index=index)
            self.models.append(mod)

    @property
    def ic(self):
        return self.base_model.ic

    @property
    def N_stars(self):
        return len(self.base_model.obs.leaves)

    @property
    def N_options(self):
        return N_options(self.N_stars, max_multiples=self.max_multiples,
                         max_stars=self.max_stars)

    @property
    def index_options(self):
        return index_options(self.N_stars)

    @property
    def model_options(self):
        return [(N, index) for N in self.N_options for index in self.index_options]

########## Utility functions ###############

def N_options(N_stars, max_multiples=1, max_stars=2):
    return [N for N in itertools.product(np.arange(max_stars) + 1, repeat=N_stars)
            if (np.array(N)>1).sum() <= max_multiples]

def index_options(N_stars):
    if N_stars==1:
        return [0]

    options = []
    for ind in itertools.product(range(N_stars), repeat=N_stars):
        diffs = np.array(ind[1:]) - np.array(ind[:-1])
        if ind[0]==0 and diffs.max()<=1:
            options.append(ind)
    return options

import numpy,import scipy,from scipy,from astropy


"""
author: ouyangtianxiong
date: 2019/12/23
des: implements attention-based emotion recognition
Based on code from https://github.com/KaihuaTang/VQA2.0-Recent-Approachs-2018.pytorch
"""
import sys
sys.path.append('../')
__author__ = 'ouyangtianxiong.bupt.edu.cn'
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.nn.utils import clip_grad_norm_
from torch.optim import Adam,SGD
from torch.nn import CrossEntropyLoss
import numpy as np
from Common_utils.model_evaluation import plot_acc_loss_curve
from Common_utils.model_training import GradualWarmupScheduler, LabelSmoothSoftmax
from Common_utils.basic_module import FCNet
import os
from data_set.seed_iv import SEED_IV, SEED_IV_DATASET

os.environ['CUDA_VISIBLE_DEVICES'] = '2'
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
class InterModalityUpdate(nn.Module):
    """
    Inter-Modality Attention Flow
    """
    def __init__(self, v_size, q_size, output_size, num_head, drop=0.0):
        super(InterModalityUpdate, self).__init__()
        self.v_size = v_size
        self.q_size = q_size
        self.output_size = output_size
        self.num_head = num_head

        self.v_lin = FCNet(v_size, output_size * 3, drop=drop)
        self.q_lin = FCNet(q_size, output_size * 3, drop=drop)

        self.v_output = FCNet(output_size + v_size, output_size, drop=drop)
        self.q_output = FCNet(output_size + q_size, output_size, drop=drop)

    def forward(self, v, q):
        """
        :param v: eeg feature [batch, regions, feature_size]
        :param q: eye feature [batch, regions, feature_size]
        :return:
        """
        batch_size, num_obj = v.shape[0], v.shape[1]
        max_len = q.shape[1]

        # transfer feature to Q, K ,V matrix, here Q, K, V are concat together
        v_tran = self.v_lin(v)
        q_tran = self.q_lin(q)
        # mask all padding object/word feature
        # split Q, K, V
        v_key, v_query, v_val = torch.split(v_tran, v_tran.size(2) // 3, dim=2)
        q_key, q_query, q_val = torch.split(q_tran, q_tran.size(2) // 3, dim=2)

        # apply multi-head
        v_key_set = torch.split(v_key, v_key.size(2) // self.num_head, dim=2)
        v_query_set = torch.split(v_query, v_query.size(2) // self.num_head, dim=2)
        v_val_set = torch.split(v_val, v_val.size(2) // self.num_head, dim=2)
        q_key_set = torch.split(q_key, q_key.size(2) // self.num_head, dim=2)
        q_query_set = torch.split(q_query, q_query.size(2) // self.num_head, dim=2)
        q_val_set = torch.split(q_val, q_val.size(2) // self.num_head, dim=2)

        # apply multi-head operation
        for i in range(self.num_head):
            v_key_slice, v_query_slice, v_val_slice = v_key_set[i], v_query_set[i], v_val_set[i]
            q_key_slice, q_query_slice, q_val_slice = q_key_set[i], q_query_set[i], q_val_set[i]
            # calculating attention
            # [batch, num_obj, max_len]
            #print(v_query_slice.shape, q_key_slice.shape)
            q2v = (v_query_slice @ q_key_slice.transpose(1, 2)) / ((self.output_size // self.num_head) ** 0.5)
            #print(q_query_slice.shape, v_key_slice.shape)
            v2q = (q_query_slice @ v_key_slice.transpose(1, 2)) / ((self.output_size // self.num_head) ** 0.5)
            # softmax attention
            interMAF_q2v = F.softmax(q2v, dim=2).unsqueeze(3) #[batch_size, num_obj, max_len, 1]
            interMAF_v2q = F.softmax(v2q, dim=2).unsqueeze(3) #[batch_size, max_len, num_obj, 1] torch.cat((v_update, (interMAF_q2v * q_val_slice.unsqueeze(1)).sum(2)), dim=2)

            #print("inter head {} \tinterMAF_q2v\n{}".format(i, interMAF_q2v.cpu().detach().numpy()))
            #print("inter head {} \tinterMAF_v2q\n{}".format(i, interMAF_v2q.cpu().detach().numpy()))
            v_update = (interMAF_q2v * q_val_slice.unsqueeze(1)).sum(2) if (i == 0) else  torch.cat((v_update, (interMAF_q2v * q_val_slice.unsqueeze(1)).sum(2)), dim=2)
            q_update = (interMAF_v2q * v_val_slice.unsqueeze(1)).sum(2) if (i == 0) else torch.cat((q_update, (interMAF_v2q * v_val_slice.unsqueeze(1)).sum(2)), dim=2)
        # update new feature
        cat_v = torch.cat((v, v_update), dim=2)
        cat_q = torch.cat((q, q_update), dim=2)
        updated_v = self.v_output(cat_v)
        updated_q = self.q_output(cat_q)
        return updated_v, updated_q

class OneSideInterModalityUpdate(nn.Module):
    """
    one-side Inter-Modality Attention Flow
    according to the paper, instead of parallel V->Q & Q->V, we first to V->Q and then Q->V
    """
    def __init__(self,src_size,tgt_size,output_size,num_head,drop=0.0):
        super(OneSideInterModalityUpdate, self).__init__()
        self.src_size = src_size
        self.tgt_size = tgt_size
        self.output_size = output_size
        self.num_head = num_head

        self.src_lin = FCNet(src_size, output_size * 2, drop=drop)
        self.tgt_lin = FCNet(tgt_size, output_size, drop=drop)

        self.tgt_output = FCNet(output_size + tgt_size, output_size, drop=drop)

    def forward(self, src, tgt):
        """
        :param src: eeg feature [batch, regions, feature_size]
        :param tgt: eye feature [batch, regions, feature_size]
        :return:
        """
        batch_size, num_src = src.shape[0],src.shape[1]
        num_tgt = tgt.shape[1]

        src_tran = self.src_lin(src)
        tgt_tran = self.tgt_lin(tgt)

        src_key, src_val = torch.split(src_tran, src_tran.size(2) // 2, dim=2)
        tgt_query = tgt_tran
        src_key_set = torch.split(src_key, src_key.size(2) // self.num_head, dim=2)
        src_val_set = torch.split(src_val, src_val.size(2) // self.num_head, dim=2)
        tgt_query_set = torch.split(tgt_query,tgt_query.size(2) // self.num_head, dim=2)
        for i in range(self.num_head):
            src_key_slice, tgt_query_slice, src_val_slice = src_key_set[i], tgt_query_set[i], src_val_set[i]
            src2tgt = (tgt_query_slice @ src_key_slice.transpose(1, 2)) / ((self.output_size // self.num_head) ** 0.5)
            interMAF_src2tgt = F.softmax(src2tgt, dim=2).unsqueeze(3)
            tgt_update = (interMAF_src2tgt * src_val_slice.unsqueeze(1)).sum(2) if (i == 0) else torch.cat((tgt_update, (interMAF_src2tgt * src_val_slice.unsqueeze(1)).sum(2)), dim=2)
        cat_tgt = torch.cat((tgt, tgt_update), dim=2)
        tgt_updated = self.tgt_output(cat_tgt)
        return tgt_updated

class DyIntraModalityUpdate(nn.Module):
    """
    Dynamic Intra-Modality Attention Flow
    """
    def __init__(self, v_size, q_size, output_size, num_head, drop=0.0):
        super(DyIntraModalityUpdate, self).__init__()
        self.v_size = v_size
        self.q_size = q_size
        self.output_size = output_size
        self.num_head = num_head

        self.v4q_gate_lin = FCNet(v_size, output_size, drop=drop)
        self.q4v_gate_lin = FCNet(q_size, output_size, drop=drop)

        self.v_lin = FCNet(v_size, output_size * 3, drop=drop)
        self.q_lin = FCNet(q_size, output_size * 3, drop=drop)

        self.v_output = FCNet(output_size, output_size, drop=drop)
        self.q_output = FCNet(output_size, output_size, drop=drop)

        self.relu = nn.ReLU()
        self.tanh = nn.Tanh()
        self.sigmoid = nn.Sigmoid()

    def forward(self, v, q):
        """
        :param v: [batch_size, num_obj, feature_size]
        :param q: [batch_size, max_len, feature_size]

        :return:
        """
        batch_size, num_obj = v.shape[0], v.shape[1]
        max_len = q.shape[1]

        # conditioned gating vector
        v_mean = v.sum(1) / num_obj
        q_mean = q.sum(1) / max_len

        v4q_gate = self.sigmoid(self.v4q_gate_lin(v_mean)).unsqueeze(1) # [batch_size, 1, feature_size]
        q4v_gate = self.sigmoid(self.q4v_gate_lin(q_mean)).unsqueeze(1) # [batch_size, 1, feature_size]

        # K, Q, V
        v_tran = self.v_lin(v)
        q_tran = self.q_lin(q)

        # split for different use
        v_key, v_query, v_val = torch.split(v_tran, v_tran.size(2) // 3, dim=2)
        q_key, q_query, q_val = torch.split(q_tran, q_tran.size(2) // 3, dim=2)

        # apply conditioned gate
        gated_v_query = (1 + q4v_gate) * v_query
        gated_v_key = (1 + q4v_gate) * v_key
        gated_v_val = (1 + q4v_gate) * v_val
        gated_q_query = (1 + v4q_gate) * q_query
        gated_q_key = (1 + v4q_gate) * q_key
        gated_q_val = (1 + v4q_gate) * q_val

        # apply multi-head
        v_key_set = torch.split(gated_v_key, gated_v_key.size(2) // self.num_head, dim=2)
        v_query_set = torch.split(gated_v_query, gated_v_query.size(2) // self.num_head, dim=2)
        v_val_set = torch.split(gated_v_val, gated_v_val.size(2) // self.num_head, dim=2)
        q_key_set = torch.split(gated_q_key, gated_q_key.size(2) // self.num_head, dim=2)
        q_query_set = torch.split(gated_q_query, gated_q_query.size(2) // self.num_head, dim=2)
        q_val_set = torch.split(gated_q_val, gated_q_val.size(2) // self.num_head, dim=2)

        for i in range(self.num_head):
            v_key_slice, v_query_slice, v_val_slice = v_key_set[i], v_query_set[i], v_val_set[i]
            q_key_slice, q_query_slice, q_val_slice = q_key_set[i], q_query_set[i], q_val_set[i]
            # calcuating attention
            v2v = (v_query_slice @ v_key_slice.transpose(1,2)) / ((self.output_size // self.num_head) ** 0.5)
            q2q = (q_query_slice @ q_key_slice.transpose(1,2)) / ((self.output_size // self.num_head) ** 0.5)
            dyIntranMAF_v2v = F.softmax(v2v, dim=2).unsqueeze(3) # [batch_size, num_obj, num_obj, 1]
            dyIntranMAF_q2q = F.softmax(q2q, dim=2).unsqueeze(3) # [batch_size, max_len, max_len, 1]
            # calculating update input
            #print("intra head {} \tinterMAF_q2v\n{}".format(i, dyIntranMAF_v2v.cpu().detach().numpy()))
            #print("intra head {} \tinterMAF_v2q\n{}".format(i, dyIntranMAF_q2q.cpu().detach().numpy()))
            v_update = (dyIntranMAF_v2v * v_val_slice.unsqueeze(1)).sum(2) if (i == 0) else torch.cat((v_update, (dyIntranMAF_v2v * v_val_slice.unsqueeze(1)).sum(2)), dim=2)
            q_update = (dyIntranMAF_q2q * q_val_slice.unsqueeze(1)).sum(2) if (i == 0) else torch.cat((q_update, (dyIntranMAF_q2q * q_val_slice.unsqueeze(1)).sum(2)), dim=2)

        # update
        updated_v = self.v_output(v + v_update)
        updated_q = self.q_output(q + q_update)
        return updated_v, updated_q

class SingleBlock(nn.Module):
    """
        Single Block Inter- and Intra modality stack multiple times, in such circumstance, all the
        basic blocks share the same parameters in the model
    """
    def __init__(self, num_blocks, v_size, q_size, output_size, num_inter_head, num_intra_head, drop=0.0):
        super(SingleBlock, self).__init__()
        self.v_size = v_size
        self.q_size = q_size
        self.output_size = output_size
        self.num_inter_head = num_inter_head
        self.num_intra_head = num_intra_head
        self.num_block = num_blocks

        self.v_lin = FCNet(v_size, output_size, drop=drop)
        self.q_lin = FCNet(q_size, output_size, drop=drop)

        self.v2q_interBlock = OneSideInterModalityUpdate(output_size, output_size, output_size, num_inter_head, drop)
        self.q2v_interBlock = OneSideInterModalityUpdate(output_size, output_size, output_size, num_inter_head, drop)
        self.intraBlock = DyIntraModalityUpdate(output_size, output_size, output_size, num_intra_head, drop)

    def forward(self, v, q):
        """
        :param v: eeg feature [batch_size, regions, feature_size]
        :param q: eye feature [batch_size, regions, feature_size]
        :return:
        """
        # transfer features
        v = self.v_lin(v)
        q = self.q_lin(q)
        # residual connection
        v_container = [v]
        q_container = [q]
        result_v = [v]
        result_q = [q]
        for i in range(self.num_block):
            q1 = self.v2q_interBlock(v_container[-1], q_container[-1])
            q_container.append(q1)
            v1 = self.q2v_interBlock(q_container[-1], v_container[-1])
            v_container.append(v1)
            v2, q2 = self.intraBlock(v_container[-1] + v_container[-2], q_container[-1] + q_container[-2])
            v_container.append(v2)
            q_container.append(q2)
            result_v.append(v1)
            result_v.append(v2)
            result_q.append(q1)
            result_q.append(q2)
            v_container.append(v_container[-1] + v_container[-2] + v_container[-3])
            q_container.append(q_container[-1] + q_container[-2] + q_container[-3])
        return sum(result_v), sum(result_q)

class MultiBlocks(nn.Module):
    """
    Stack multiple single block layer, each layer possess their own parameters
    """

    def __init__(self, num_blocks, v_size, q_size, output_size, num_inter_head, num_intra_head, drop=0.0):
        super(MultiBlocks, self).__init__()
        self.v_size = v_size
        self.q_size = q_size
        self.output_size = output_size
        self.num_inter_head = num_inter_head
        self.num_intra_head = num_intra_head
        self.num_blocks = num_blocks

        self.v_lin = FCNet(v_size, output_size, drop=drop)
        self.q_lin = FCNet(q_size, output_size, drop=drop)

        blocks = []
        for i in range(self.num_blocks):
            #blocks.append(OneSideInterModalityUpdate(output_size, output_size, output_size, num_inter_head, drop))
            #blocks.append(OneSideInterModalityUpdate(output_size, output_size, output_size, num_inter_head, drop))
            blocks.append(InterModalityUpdate(output_size, output_size, output_size, num_inter_head, drop))
            blocks.append(DyIntraModalityUpdate(output_size, output_size, output_size, num_intra_head, drop))
        self.multi_blocks = nn.ModuleList(blocks)

    def forward(self, v, q):
        """
        :param v: eeg feature [batch, regions, feature_size]
        :param q: eye feature [batch, regions, feature_size]
        :return:
        """
        # transfer feature
        v = self.v_lin(v)
        q = self.q_lin(q)
        v_container = [v]
        q_container = [q]
        result_v = [v]
        result_q = [q]

        # dense residule connection
        for i in range(self.num_blocks):
            # q1 = self.multi_blocks[i * 3 + 0](v_container[-1], q_container[-1])
            # q_container.append(q1)
            # v1 = self.multi_blocks[i * 3 + 1](q_container[-1], v_container[-1])
            # v_container.append(v1)
            v1, q1 = self.multi_blocks[i * 2 + 0](v_container[-1], q_container[-1])
            q_container.append(q1)
            v_container.append(v1)
            v2, q2 = self.multi_blocks[i * 2 + 1](v_container[-1] + v_container[-2], q_container[-1] + q_container[-2])
            v_container.append(v2)
            q_container.append(q2)
            result_v.append(v1)
            result_v.append(v2)
            result_q.append(q1)
            result_q.append(q2)
            v_container.append(v_container[-1] + v_container[-2] + v_container[-3])
            q_container.append(q_container[-1] + q_container[-2] + q_container[-3])
        return sum(result_v), sum(result_q)

class EEGFeatureExtractor(nn.Module):
    def __init__(self, eeg_size, output_size):
        super(EEGFeatureExtractor, self).__init__()
        self.eeg_size = eeg_size
        self.output_size = output_size
        self.regions = 16  # regions的数量
        self.regions_indexs = [torch.LongTensor(e) for e in
                               [[3, 0, 1, 2, 4], [7, 8, 9, 10, 11], [5, 6], [13, 12], [14, 15, 23, 24, 32, 33],
                                [22, 21, 31, 30, 40, 39], [16, 17, 18, 19, 20], [25, 26, 27, 28, 29],
                                [34, 35, 36, 37, 38], [41, 42], [49, 48], [43, 44, 45, 46, 47],
                                [50, 51, 57], [56, 55, 61], [52, 53, 54], [58, 59, 60]]]
        reginal_extractors = []
        for i in range(self.regions):
            reginal_extractors.append(nn.LSTM(input_size=eeg_size, hidden_size= output_size // 2, batch_first=True, bias=True, bidirectional=True))

        self.reginalFeatureExtractors = nn.ModuleList(reginal_extractors)
        self.bn = nn.BatchNorm1d(num_features=self.regions)

    def forward(self, x):
        """
        :param x: [batch, n_electrode, 5]
        :return: [batch, regions, feature_size]
        """
        batch, n_electrode, _ = x.shape
        X_regions_input = []  # 列表存储不同区域的张量输入
        for i in range(self.regions):
            X_regions_input.append(x.index_select(dim=1, index=self.regions_indexs[i].to(device)))
        X_regional_lstm_out = []
        for i in range(self.regions):
            shape = X_regions_input[i].shape
            # print(shape)
            # 先转成（B*T,n_i,d）再进LSTM
            hidden_units, _ = self.reginalFeatureExtractors[i](X_regions_input[i].reshape((-1, shape[-2], shape[-1])))
            X_regional_lstm_out.append(hidden_units[:, -1, :].squeeze())
        # X_regional_feature : 列表：元素为tensor [ B*T, regions_num, 2*self.d_r]
        # reshape成(B*T, regions, 2*self.d_r)
        # (B * T, regions, 2* self.d_r)
        X_regional_feature = torch.cat(X_regional_lstm_out, dim=-1).reshape(batch, self.regions, self.output_size)
        return self.bn(X_regional_feature)

class EYEFeatureExtractor(nn.Module):
    def __init__(self, eye_size, output_size):
        super(EYEFeatureExtractor, self).__init__()
        self.eye_size = eye_size
        self.output_size = output_size
        self.regions = 5
        self.regions_indexs = [torch.LongTensor(e) for e in
                               [[0,1,2,3,4,5,6,7,8,9,10,11],
                                [12,13,14,15],
                                [16,17],
                                [18,19,20,21],
                                [22,23,24,25,26,27,28,29,30]]]

        eye_extractor = []
        eye_extractor.append(FCNet(in_size=12, out_size=output_size, activate='relu'))
        eye_extractor.append(FCNet(in_size=4, out_size=output_size, activate='relu'))
        eye_extractor.append(FCNet(in_size=2, out_size=output_size, activate='relu'))
        eye_extractor.append(FCNet(in_size=4, out_size=output_size, activate='relu'))
        eye_extractor.append(FCNet(in_size=9, out_size=output_size, activate='relu'))
        self.eyeFeatureExtractor = nn.ModuleList(eye_extractor)
        self.bn = nn.BatchNorm1d(num_features=self.regions)

    def forward(self, x):
        """
        :param x: EYE feature [batch, 31]
        :return: [batch, regons, output_size]
        """
        B = x.shape[0]
        X_regional_output = []
        for i in range(self.regions):
            X_regional_output.append(self.eyeFeatureExtractor[i](x.index_select(dim=1, index=self.regions_indexs[i].to(device))))
        X_regional_feature = torch.cat(X_regional_output, dim=-1).reshape(B, self.regions, self.output_size)
        return self.bn(X_regional_feature)

class Classifier(nn.Sequential):
    def __init__(self, in_features, mid_features, out_features, drop=0.0):
        super(Classifier, self).__init__()
        # define number of detector for each sentiment class
        self.lin1 = FCNet(in_features, mid_features, activate='relu', drop=drop)
        self.lin2 = FCNet(mid_features, out_features, drop=drop)
        #
        self.bilinear = nn.Bilinear(in1_features=in_features, in2_features=in_features, out_features=in_features)
    def forward(self, v, q):
        """
        :param v: [batch, r1, features]
        :param q: [batch, r2, features]
        :return:
        """
        num_obj = v.shape[2]
        max_len = q.shape[2]

        v_mean = v.sum(1) / num_obj
        q_mean = q.sum(1) / max_len
        #print("classifier v_mean", v_mean[0])
        #print("classifier q_mean", q_mean[0])

        #out = self.lin1(v_mean * q_mean)
        out = self.lin1(self.bilinear(v_mean, q_mean))
        #print("classifier out 1", out[0])
        out = self.lin2(out)
        #print("classifier out 2", out[0])
        return out
class Hierarchical_ATTN(nn.Module):
    def __init__(self):
        super(Hierarchical_ATTN, self).__init__()
        self.eye_features = 16 # 256
        self.eeg_features = 16 # 256
        self.hidden_feature = 32 # 256
        self.num_inter_head = 4
        self.num_intra_head = 4
        self.num_block = 1

        assert self.hidden_feature % self.num_inter_head == 0, 'hidden features size can not be divided by header nums, please check!!'
        assert self.hidden_feature % self.num_inter_head == 0, 'hidden features size can not be divided by header nums, please check!!'

        # basic feature extractor
        self.eegFeatureExtractor = EEGFeatureExtractor(eeg_size=5, output_size=self.eeg_features)

        self.eyeFeatureExtractor = EYEFeatureExtractor(eye_size=31, output_size=self.eeg_features)

        # inter- & intra-modality attention flow mechanism for fusion cross modality feature
        self.interIntraBlocks = MultiBlocks(
            num_blocks=self.num_block,
            v_size=self.eeg_features,
            q_size=self.eye_features,
            output_size=self.hidden_feature,
            num_inter_head=self.num_inter_head,
            num_intra_head=self.num_intra_head,
            drop=0.1
        )

        # emotion classifier
        self.classifier = Classifier(
            in_features=self.hidden_feature,
            mid_features=512, out_features=4,
            drop=0.5)
        # self.classifier = Senti_Map_Classifier(
        #     in_features=self.hidden_feature,
        #     mid_features=256, out_features=4,
        #     drop=0.5)
    def forward(self, v, q):
        """
        :param v: eeg feature [batch, n, 5]
        :param q:  eye feature [batch, 31]
        :return: predict logits [batch, max_answer]
        """
        # prepare v & q feature

        v = self.eegFeatureExtractor(v)
        q = self.eyeFeatureExtractor(q)

        # feature normalization
        v = v / (v.norm(p=2, dim=2, keepdim=True) + 1e-12).expand_as(v) # [batch, num_obj, feature]
        q = q / (q.norm(p=2, dim=2, keepdim=True) + 1e-12).expand_as(q)

        # inter- & intra- modality attention flow
        v, q = self.interIntraBlocks(v, q)

        # predict logits
        answer = self.classifier(v, q)
        return answer

def main(session=1, mode='subject_dependent'):
    os.environ['CUDA_VISIBLE_DEVICES'] = '2'
    # prepare data
    session = session
    balance = False
    shuffle = False
    modal = 'concat'
    nor_method = 1
    label_smooth = 0.1
    fine_tuning = True

    # reading the data in the whole dataset
    all_individual_data = []
    for i in range(1, 16):
        print("contructing dataset...")
        eeg = SEED_IV(session=session, individual=i, modal=modal, shuffle=shuffle, balance=balance,
                      normalization=nor_method)
        _train_X, _train_Y = eeg.get_train_data()
        _test_X, _test_Y = eeg.get_test_data()
        all_individual_data.append([(_train_X, _train_Y), (_test_X, _test_Y)])

    # Hyper-parameters
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    epochs = 120
    batch_size = 128
    learning_rate = 1e-3
    criterion = LabelSmoothSoftmax(lb_smooth=label_smooth)
    for idx in range(1, 16):
        if mode == 'subject_dependent':
            train_X, train_Y = all_individual_data[idx-1][0][0], all_individual_data[idx-1][0][1]
            test_X, test_Y = all_individual_data[idx-1][1][0], all_individual_data[idx-1][1][1]
            exp_des = "%d_dependent_in_seesion_%d_%s_%s_%s_%d_%d" % (
                idx, session, 'balance' if balance else 'without_balance',
                'shuffle' if shuffle else "without_shuffle", 'seed', epochs, batch_size)
            print("starting subject-dependent training experiments on individual %d in session %d"% (idx, session))
        elif mode == 'subject_independent':
            train_X = np.vstack([np.vstack((e[0][0], e[1][0])) for i, e in enumerate(all_individual_data) if i != idx-1])
            train_Y = np.hstack([np.hstack((e[0][1], e[1][1])) for i, e in enumerate(all_individual_data) if i != idx-1])
            test_X = np.vstack((all_individual_data[idx-1][0][0], all_individual_data[idx-1][1][0]))
            test_Y = np.hstack((all_individual_data[idx-1][0][1], all_individual_data[idx-1][1][1]))
            exp_des = "%d_independent_as_testset_in_seesion_%d_%s_%s_%s_%d_%d" % (
                idx, session, 'balance' if balance else 'without_balance',
                'shuffle' if shuffle else "without_shuffle", 'seed', epochs, batch_size)
            print("starting subject-independent training experiments with individual %d in session %d as test set" % (idx, session))
        else:
            raise ValueError

        print("train_X shape", train_X.shape)
        print("train_Y shape", train_Y.shape)
        print("test_X shape", test_X.shape)
        print("test_Y shape", test_Y.shape)
        train_loader = DataLoader(dataset=SEED_IV_DATASET(train_X, train_Y), batch_size=batch_size, shuffle=shuffle,
                                  num_workers=4)
        test_loader = DataLoader(dataset=SEED_IV_DATASET(test_X, test_Y), batch_size=batch_size, shuffle=shuffle,
                                 num_workers=4)

        print("model construction...")
        net = Hierarchical_ATTN()
        # if fine_tuning we continue train the pretrained model
        if mode == 'subject_dependent' and fine_tuning:
            load_path = "../../saved_models/%s/session_%d/subject_%d_as_testset" % (net.__class__.__name__, session, idx)
            files = os.listdir(load_path)
            best_model = max(files)
            checkpoint = torch.load(os.path.join(load_path, best_model))
            net.load_state_dict(checkpoint['net'])
            learning_rate = 1e-5
            batch_size = train_X.shape[0]

        net = net.to(device)
        save_model_path = '../../saved_models/%s/session_%d/subject_%d_as_testset' % (
        net.__class__.__name__, session, idx) if mode == 'subject_independent' else '../../saved_models/%s/session_%d/subject_%d' % (
        net.__class__.__name__, session, idx)
        if not os.path.exists(save_model_path):
            os.makedirs(save_model_path)
        optimization = Adam(net.parameters(), lr=learning_rate, weight_decay=0.001)

        # save model training state
        running_loss_list = []
        running_acc_list = []
        testing_loss_list =[]
        testing_acc_list = []
        best_acc = -1
        print("start training...")
        scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer=optimization, T_max=epochs)
        scheduler_warmup = GradualWarmupScheduler(optimizer=optimization, multiplier=10,
                                                        total_epoch=np.ceil(0.1 * epochs),
                                                        after_scheduler=scheduler_cosine)
        for epoch in range(epochs):
            net.train()
            running_loss = 0.0
            correct = 0.0
            total = 0.0
            for i, (feature, target) in enumerate(train_loader):
                optimization.zero_grad()
                #print("脏数据统计", torch.sum(torch.isnan(feature), dim=0))
                eeg = feature[:, :310]
                eye = feature[:, 310:]
                eeg = eeg.reshape(-1, 62, 5)
                eeg = eeg.to(device)
                eye = eye.to(device)
                target = target.type(torch.LongTensor).to(device)
                out = net(eeg, eye)
                #print("batch output",out[0])
                cross_entropy_loss = criterion(out, target)
                cross_entropy_loss.backward()
                clip_grad_norm_(net.parameters(), max_norm=10)
                optimization.step()
                running_loss += cross_entropy_loss.item()
                #print("batch loss", loss.item())
                _, prediction = torch.max(out.data, dim=-1)
                total += target.size(0)
                correct += prediction.eq(target.data).cpu().sum()
            cur_loss = running_loss / len(train_loader)
            cur_acc = correct / total
            if isinstance(cur_acc, torch.Tensor):
                cur_acc = cur_acc.item()
            if isinstance(cur_loss, torch.Tensor):
                cur_loss = cur_loss.item()
            print('Loss: %.10f | Acc: %.3f%% (%d/%d)' % (
                cur_loss, 100 * cur_acc, correct, total))
            running_loss_list.append(cur_loss)
            running_acc_list.append(cur_acc)
            scheduler_warmup.step()
            if epoch % 1 == 0:
                net.eval()
                print("start evaluating...")
                testing_loss = 0.0
                test_correct = 0.0
                test_total = 0.0
                for i, (feature, target) in enumerate(test_loader):
                    eeg = feature[:, :310]
                    eye = feature[:, 310:]
                    eeg = eeg.reshape(-1, 62, 5)
                    eeg = eeg.to(device)
                    eye = eye.to(device)
                    target = target.type(torch.LongTensor).to(device)
                    with torch.no_grad():
                        out = net(eeg, eye)
                        loss = criterion(out, target)
                        testing_loss += loss.item()
                        _, prediction = torch.max(out.data, dim=-1)
                        # print(prediction)
                        test_total += target.size(0)
                        test_correct += prediction.eq(target.data).cpu().sum()
                test_acc = test_correct / test_total
                test_loss = testing_loss / len(test_loader)
                if isinstance(test_acc, torch.Tensor):
                    test_acc = test_acc.item()
                if isinstance(test_loss, torch.Tensor):
                    test_loss = test_loss.item()
                print('Testset Loss: %.10f | Acc: %.3f%% (%d/%d)' % (
                    test_loss, 100 * test_acc, test_correct, test_total))
                testing_acc_list.append(test_acc)
                testing_loss_list.append(test_loss)
                if test_acc > best_acc:
                    best_acc = test_acc
                    print("better model founded in testsets, start saving new model")
                    model_name = '%s_%s' % (net.__class__.__name__, str(best_acc)[2:6])
                    state = {
                        'net': net.state_dict(),
                        'epoch': epoch,
                        'best_acc': best_acc,
                        'current_loss': test_loss
                    }
                    torch.save(state, os.path.join(save_model_path, model_name))
        plot_acc_loss_curve({'train_loss': running_loss_list,
                            'train_acc': running_acc_list,
                            'test_loss': testing_loss_list,
                            'test_acc': testing_acc_list}, net.__class__.__name__, exp_des)
if __name__ == '__main__':
    for mode in ['subject_independent','subject_dependent']:
        for session in range(1, 4):
            main(session, mode)
    # main(1, 'subject_dependent')
    print("experiment done!")