import os
import json
import numpy as np
import pandas as pd

from pathlib import Path

from .half_sarcomere import half_sarcomere as hs
from .heart_rate import heart_rate as hr
from .baroreflex import baroreflex as br
from .growth import growth as gr
from .vad import vad as va

from protocol import protocol as prot
from sim_options import sim_options as sim_opt

class single_ventricle_circulation():
    """Class for a single ventricle circulation"""

    # Import functions from additional files
    from .calculate_properties import \
        return_wall_thickness, \
        return_lv_circumference, \
        return_internal_radius_for_chamber_volume, \
        return_lv_pressure, \
        return_stroke_work

    from .write_data import \
        write_complete_data_to_sim_data, \
        write_complete_data_to_envelope_data, \
        write_envelope_data_to_sim_data, \
        write_output_files, \
        write_sim_results_to_file,\
        run_output_handler

    def __init__(self, model_json_file_string, thread_id=[]):

        # Check for file
        if (model_json_file_string == []):
            print('No model file specified. Cannot create model')
            return

        # Set thread id
        self.thread_id = thread_id

        # Status
        print('Initialising single_ventricle_circulation from %s on thread_id %i'
              % (model_json_file_string, self.thread_id))

        # Load the model as a dict
        with open(model_json_file_string, 'r') as f:
            sc_model = json.load(f)

        # Create a model dict for things that do not change during a simulation
        self.model = dict()
        # And a data dict for things that might
        self.data = dict()

        # Initialize the circulation object using data from the json file
        circ_model = sc_model["circulation"]

        # Store the number of compartments
        self.model['no_of_compartments'] = len(circ_model['compartments'])

        # Initialise the time
        self.data['time'] = 0

        # Store blood volume
        self.data['blood_volume'] = circ_model['blood_volume']

        vessels_list = []
        for comp in circ_model['compartments']:
            if not (comp['name'] == 'ventricle'):
                n = comp['name']
                vessels_list.append(comp['name'])
                for t in ['resistance', 'compliance', 'slack_volume',
                          'inertance']:
                    n = ('%s_%s') % (comp['name'], t)
                    self.data[n] = comp[t]
            else:
                # Ventricle
                self.data['ventricle_resistance'] = comp['resistance']
                self.data['ventricle_slack_volume'] = comp['slack_volume']
                self.data['ventricle_wall_volume'] = comp['wall_volume']
                self.model['ventricle_wall_density'] = comp['wall_density']
                self.model['ventricle_inertance'] = comp['inertance']

        # Build the compliance, inertance, and resistance arrays
        self.data['compliance'] = []
        for v in vessels_list:
            c = self.data[('%s_compliance' % v)]
            self.data['compliance'].append(c)
        # Add in 0 for ventricular compliance
        self.data['compliance'].append(0)
        # Convert to numpy array
        self.data['compliance'] = np.array(self.data['compliance'])

        self.data['inertance'] = []
        for v in vessels_list:
            i = self.data[('%s_inertance' % v)]
            self.data['inertance'].append(i)
        # Add in 0 for ventricular inertance
        self.data['inertance'].append(0)
        # Convert to numpy array
        self.data['inertance'] = np.array(self.data['inertance'])

        self.data['resistance'] = []
        vessels_list.append('ventricle')
        self.model['compartment_list'] = vessels_list
        for v in self.model['compartment_list']:
            r = self.data[('%s_resistance' % v)]
            self.data['resistance'].append(r)
        self.data['resistance'] = np.array(self.data['resistance'])

        # Add in conductances for valves
        self.data['mitral_insufficiency_conductance'] = 0
        self.data['aortic_insufficiency_conductance'] = 0

        # Create a heart-rate object
        self.hr = hr.heart_rate(sc_model['heart_rate'])

        # Pull off the half_sarcomere parameters and make a half-sarcomere
        self.hs = hs.half_sarcomere(sc_model['half_sarcomere'], self)

        # Deduce the hsl where force is zero and set the hsl to that length
        self.hs.data['slack_hsl'] = \
            self.hs.myof.return_hs_length_for_stress(0.0)
        delta_hsl = self.hs.data['slack_hsl'] - self.hs.data['hs_length']
        self.hs.update_simulation(0.0, delta_hsl, 0.0)

        # Deduce the slack circumference of the ventricle and deduce
        # the number of sarcomeres by dividing the slack_hsl
        self.data['ventricle_circumference'] = \
            self.return_lv_circumference(self.data['ventricle_slack_volume'])
        self.data['n_hs'] = 1e9*self.data['ventricle_circumference'] / \
            self.hs.data['slack_hsl']

        # Create and fill arrays for the volume, slack_volume and pressure
        self.data['v'] = np.zeros(self.model['no_of_compartments'])
        self.data['s'] = np.zeros(self.model['no_of_compartments'])
        # Put most of the blood in the veins
        for i, c in enumerate(self.model['compartment_list']):
            n = ('%s_slack_volume' % c)
            self.data['s'][i] = self.data[n]
            self.data['v'][i] = self.data[n]
        self.data['total_slack_volume'] = sum(self.data['s'])

        # Excess blood goes in veins
        self.data['v'][-2] = self.data['v'][-2] + \
            (self.data['blood_volume'] - self.data['total_slack_volume'])

        # Set the thick wall multiplier
        # 0 for thin_wall, 1 for thick_wall
        self.thick_wall_multiplier = 0
        if (sc_model['half_sarcomere']['myofilaments']['implementation']['thick_wall_approximation']):
            self.thick_wall_multiplier = 1

        self.data['p'] = np.zeros(self.model['no_of_compartments'])
        for i in np.arange(0, self.model['no_of_compartments']-1):
            self.data['p'][i] = (self.data['v'][i] - self.data['s'][i]) / \
                self.data['compliance'][i]
        self.data['p'][-1] = self.return_lv_pressure(self.data['v'][-1])

        # Allocate space for pressure, volume and slack_volume
        for i, v in enumerate(self.model['compartment_list']):
            self.data['pressure_%s' % v] = self.data['p'][i]
            self.data['volume_%s' % v] = self.data['v'][i]
            self.data['slack_volume_%s' % v] = self.data['s'][i]

        # Add in the wall thickness
        self.data['ventricle_wall_thickness'] = self.return_wall_thickness(
            self.data['volume_ventricle'])

        # Allocate space for flows
        self.data['f'] = np.zeros(self.model['no_of_compartments'])
        if (self.model['no_of_compartments'] == 7):
            self.model['flow_list'] = \
                ['flow_ventricle_to_aorta',
                 'flow_aorta_to_arteries',
                 'flow_arteries_to_arterioles',
                 'flow_arterioles_to_capillaries',
                 'flow_capillaries_to_venules',
                 'flow_venules_to_veins',
                 'flow_veins_to_ventricle',
                 'flow_VAD']
            for f in self.model['flow_list']:
                self.data[f] = 0
        else:
            print('Initialisation: Flows need to be rebuilt for model')
            exit(1)

        # Set the heart rate
        self.data['heart_rate'] = self.hr.return_heart_rate()

        # If requried, create the baroreceptor
        self.data['baro_active'] = 0
        if ('baroreflex' in sc_model):
            self.br = br.baroreflex(sc_model['baroreflex'],
                                    self,
                                    self.data['pressure_arteries'])
        else:
            self.br = []

        # Create the growth modules
        self.data['growth_active'] = 0
        self.data['growth_dn'] = 0
        self.data['growth_dm'] = 0

        if ('growth' in sc_model):
            self.gr = gr.growth(sc_model['growth'], self)
        else:
            self.gr = []

        # Add functional fields
        self.data['stroke_volume'] = 0
        self.data['ejection_fraction'] = 0
        self.data['stroke_work'] = 0
        self.data['stroke_cost'] = 0
        self.data['efficiency'] = 0
        self.data['mean_arterial_pressure'] = 0

        # Set the last index for heart_beat initiation
        self.last_heart_beat_time = -1

        # If required create a VAD
        if ('VAD' in sc_model):
            self.va = va.VAD(sc_model['VAD'], self)
        else:
            self.va = []

        # Save space for some filenames
        self.sim_results_file_string = []
        self.output_handerl_file_string = []

        # Create a sim_options object
        self.so = []

    def create_data_structure(self, no_of_data_points):
        """ returns a data frame from the data dicts of each component """

        sim_data = pd.DataFrame()
        z = np.zeros(no_of_data_points)

        # Prune some fields from the self_data
        sk = []
        for k in self.data.keys():
            if (k not in ['p','v','s','compliance','resistance',
                          'inertance','f']):
                sk.append(k)

        data_fields = sk + \
            list(self.hr.data.keys()) + \
            list(self.hs.data.keys()) + \
            list(self.hs.memb.data.keys()) + \
            list(self.hs.myof.data.keys())
        if hasattr(self.hs, 'ener'):
            data_fields = data_fields + \
                list(self.hs.ener.data.keys())
        data_fields = data_fields + ['write_mode']

        # Add in fields from optional modules
        if (self.br != []):
            data_fields = data_fields + list(self.br.data.keys())
        if (self.gr != [] ):
            data_fields = data_fields + list(self.gr.data.keys())
        if (self.va != []):
            data_fields = data_fields + list(self.va.data.keys())

        for f in data_fields:
            s = pd.Series(data=z, name=f)
            sim_data = pd.concat([sim_data, s], axis=1)

        return sim_data

    def run_simulation(self,
                       protocol_file_string=[],
                       output_handler_file_string=[],
                       sim_options_file_string=[],
                       sim_results_file_string=[]):
        """ Run the simulation """

        # Load the protocol
        if (protocol_file_string == []):
            print("No protocol_file_string. Exiting")
            return
        self.prot = prot.protocol(protocol_file_string)

        # Check for sim_options
        if (sim_options_file_string):
            self.so = sim_opt.sim_options(sim_options_file_string,
                                          self.prot.data['time_step'],
                                          self)
        else:
            self.so = []

        # Determine the number of data points in the output file
        # If burst mode has been set in the sim_options, we need to create
        # two data structures - one for the main data, and the other that
        # can be used to calculate the min and max values in a given time
        # envelope
        # Otherwise, we just need one structure with a row for every
        # time-point in the simulation
        if ('n_burst_points' in self.so.data):
            self.sim_data = \
                self.create_data_structure(self.so.data['n_burst_points'])
            self.envelope_data = \
                self.create_data_structure(self.so.data['n_envelope_points'])
        else:
            self.sim_data = \
                self.create_data_structure(self.prot.data['no_of_time_steps'])

        # Update the file handles
        self.sim_results_file_string = sim_results_file_string
        self.output_handler_file_string = output_handler_file_string

        # Step through the simulation
        time_step = self.prot.data['time_step']
        self.t_counter = 0
        self.write_counter = 0
        self.envelope_counter = 0
        for i in np.arange(self.prot.data['no_of_time_steps']):
            # Iplmement a step
            self.implement_time_step(time_step)
            # Check whether we should dump the output
            if ('periodic_save_interval_s' in self.so.data):
                if ((i > 0) and
                    (np.mod(self.data['time'],
                            self.so.data['periodic_save_interval_s'])
                                 < time_step) and
                    (i < (self.prot.data['no_of_time_steps']-1))):
                    # Write sim data to file and run the output handler
                    self.write_output_files()

        # Tidy up by writing sim data to file and running the output handler
        self.write_output_files()

    def return_system_values(self, time_interval=3):
        d = dict()
        if (self.data['time'] > time_interval):
            self.temp_data = \
                self.sim_data[self.sim_data['time'].between(
                    self.data['time']-time_interval, self.data['time'])]
            d['volume_ventricle_max'] = \
                self.temp_data['volume_ventricle'].max()
            d['stroke_volume'] = d['volume_ventricle_max'] - \
                self.temp_data['volume_ventricle'].min()
            d['ejection_fraction'] = self.temp_data['ejection_fraction'].mean()
            d['heart_rate'] = self.data['heart_rate']
            d['cardiac_output'] = d['stroke_volume'] * d['heart_rate']
            d['hs_length_max'] = self.temp_data['hs_length'].max()
            d['fractional_shortening'] = \
                (d['hs_length_max'] - self.temp_data['hs_length'].min()) / \
                d['hs_length_max']
            d['cpt_int_pas_stress_mean'] = self.temp_data['cpt_int_pas_stress'].mean()
            d['cpt_cb_stress_mean'] = self.temp_data['cpt_cb_stress'].mean()
            d['efficiency'] = self.data['efficiency']
            d['pressure_artery_max'] = \
                self.temp_data['pressure_arteries'].max()
            d['pressure_artery_min'] = \
                self.temp_data['pressure_arteries'].min()
            d['pressure_veins_mean'] = \
                self.temp_data['pressure_veins'].mean()
            d['volume_veins_proportion'] = \
                self.temp_data['volume_veins'].mean() / \
                self.data['blood_volume']
            d['n_hs'] = self.data['n_hs']
            d['dn'] = self.data['growth_dn']
            d['ventricle_wall_volume'] = self.data['ventricle_wall_volume']
            d['dm'] = self.data['growth_dm']

            if ('ener_intracell_ATP_conc' in self.temp_data):
                d['ener_intracell_ATP_conc'] = \
                    self.temp_data['ener_intracell_ATP_conc'].mean()

            if ('baro_B_a' in self.temp_data):
                d['baro_B_a'] = self.temp_data['baro_B_a'].mean()
            if ('baro_B_b' in self.temp_data):
                d['baro_B_b'] = self.temp_data['baro_B_b'].mean()
            if ('baro_B_c_heart_rate_t_quiescent_period' in self.temp_data):
                d['baro_B_c_heart_rate'] = self.temp_data['baro_B_c_heart_rate_t_quiescent_period'].mean()

        return d

    def implement_time_step(self, time_step):
        """ Implements time step """
        self.data['time'] = self.data['time'] + time_step

        # Display progress
        if (self.t_counter % 1000 == 0):
            print('Thread_id[%i]: Sim time (s): %.0f  %.0f%% complete' %
                  (self.thread_id, self.data['time'],
                   100*self.t_counter/self.prot.data['no_of_time_steps']))
            system_values = self.return_system_values()
            print(json.dumps(system_values, indent=4))

        # Check for baroreflex and implement
        if (self.br):
            self.data['baro_active'] = 0
            for b in self.prot.baro_activations:
                if ((self.t_counter >= b.data['t_start_ind']) and
                        (self.t_counter < b.data['t_stop_ind'])):
                    self.data['baro_active'] = 1

            self.br.implement_time_step(self.data['pressure_arteries'],
                                        time_step,
                                        reflex_active=
                                        self.data['baro_active'])

        # Check and implement perturbations
        for p in self.prot.perturbations:

            if((self.t_counter >= p.data['t_start_ind']) and
               (self.t_counter < p.data['t_stop_ind'])):
                if (p.data['level'] == 'myofilaments'):
                    self.hs.myof.data[p.data['variable']] += \
                        p.data['increment']
                elif (p.data['level'] == 'VAD'):
                    self.va.data[p.data['variable']] += \
                        p.data['increment']
                elif (p.data['level'] == 'circulation'):
                    self.data[p.data['variable']] += \
                        p.data['increment']
                elif (p.data['level'] == 'baroreflex'):
                    self.br.data[p.data['variable']] += \
                        p.data['increment']
                elif (p.data['level'] == 'growth'):
                    self.gr.data[p.data['variable']] += \
                        p.data['growth']

                self.check_baroreflex_perturbations(p)

        # Rebuild system arrays
        self.rebuild_from_perturbations()

        # Check for growth module and implement
        if (self.gr):
            self.data['growth_active'] = 0
            for g in self.prot.growth_activations:
                if ((self.t_counter >= g.data['t_start_ind']) and
                        (self.t_counter < g.data['t_stop_ind'])):
                    self.data['growth_active'] = 1

            self.gr.implement_time_step(time_step,
                                        growth_active=
                                        self.data['growth_active'])

        # Run the hr module
        (activation, new_beat) = self.hr.implement_time_step(time_step)

        # Advance half_sarcomere forward in time
        # First update the kinetic steps
        self.hs.update_simulation(time_step, 0, activation)

        # Now evolve volumes
        self.data['v'] = self.evolve_volumes(time_step, self.data['v'])

        # Apply the length change to the half-sarcomere
        new_circumference = self.return_lv_circumference(
                                self.data['v'][-1])
        delta_circumference = new_circumference - \
            self.data['ventricle_circumference']
        delta_hsl = ((1e9*delta_circumference) -
                     (self.hs.data['hs_length'] * self.data['growth_dn'])) / \
            self.data['n_hs']

        # Update model
        self.data['ventricle_circumference'] = new_circumference
        self.data['ventricle_wall_volume'] = \
            self.data['ventricle_wall_volume'] + \
            self.data['growth_dm'] + \
            (self.data['ventricle_wall_volume'] *
             self.data['growth_dn'] / self.data['n_hs'])
        self.data['n_hs'] = self.data['n_hs'] + self.data['growth_dn']

        # Update the half-sarcomere with the new length
        self.hs.update_simulation(0, delta_hsl, 0)

        # Update the pressures
        for i in range(self.model['no_of_compartments']-1):
            self.data['p'][i] = (self.data['v'][i] - self.data['s'][i]) / \
                self.data['compliance'][i]
        self.data['p'][-1] = self.return_lv_pressure(self.data['v'][-1])

        # Update the objects' data
        self.update_data(time_step, new_beat)
        self.hs.update_data(new_beat)

        # Now that ventricular volume is calculated, update the wall thickness
        self.data['ventricle_wall_thickness'] = self.return_wall_thickness(
            self.data['volume_ventricle'])

        # Now update the sim_data if appropriate
        # This requires checking if we are in burst mode and if so,
        # what write mode we are in
        save_mode = 1
        burst_mode = 'complete'
        if ('n_burst_points' in self.so.data):
            # Work out what we want to write to file
            (save_mode, burst_mode) = \
                self.so.return_save_status(self.data['time'])
            # Update the envelope data
            self.write_complete_data_to_envelope_data(self.envelope_counter)

        if (save_mode == 1):
            # Write full data to sim_data
            self.write_complete_data_to_sim_data(self.write_counter)
        if (save_mode == 2):
            # Write envelope data to sim_data
            self.write_envelope_data_to_sim_data(self.write_counter)

        # Dump cb distributions if required
        if self.so:
            if ('cb_dump_file_string' in self.so.data):
                if ((self.t_counter >=
                     self.so.data['cb_dump_t_start_ind']) and
                    (self.t_counter <
                     self.so.data['cb_dump_t_stop_ind'])):
                    self.so.append_cb_distribution(self.data['time'])

        # Update the last heart beat time if required
        if (new_beat > 0):
            self.last_heart_beat_time = self.data['time']

        # Update the t counter for the next step
        self.t_counter = self.t_counter + 1

    def check_baroreflex_perturbations(self, p):
        """ Checks whether a parameter under reflex control has been
            externally perturned and adjusts values accordingly """

        for bc in self.br.controls:
            if (bc.data['variable'] == p.data['variable']):
                bc.data['basal_value'] = bc.data['basal_value'] + \
                    p.data['increment']
                bc.data['para_value'] == bc.data['para_factor'] * \
                    bc.data['basal_value']
                bc.data['symp_value'] == bc.data['symp_factor'] * \
                    bc.data['basal_value']

    def return_min_max(self, data_frame):
        """ returns list of min and max values from a data frame """
        min_value = data_frame.min()
        max_value = data_frame.max()

        return min_value, max_value

    def rebuild_from_perturbations(self):
        """ builds system arrays that could change during simulation """

        for i, v in enumerate(self.model['compartment_list']):
            r = self.data[('%s_resistance' % v)]
            self.data['resistance'][i] = r

        for i, v in enumerate(self.model['compartment_list']):
            if (i < (self.model['no_of_compartments']-1)):
                c = self.data[('%s_compliance' % v)]
                self.data['compliance'][i] = c

    def update_data(self, time_step, new_beat):
        """ Update data after a time step """

        # Update data for the heart-rate
        self.data['heart_rate'] = self.hr.return_heart_rate()

        self.data['f'] = self.return_flows(self.data['v'], time_step)
        for i, f in enumerate(self.model['flow_list']):
            self.data[f] = self.data['f'][i]

        for i, v in enumerate(self.model['compartment_list']):
            self.data['pressure_%s' % v] = self.data['p'][i]
            self.data['volume_%s' % v] = self.data['v'][i]

        if (hasattr(self.hs, 'ener')):
            self.data['ener_ATPase_to_myo'] = \
                self.hs.ener.data['ener_flux_ATP_consumed'] / \
                    (0.001 * self.data['ventricle_wall_volume'] *
                     (1.0 - self.hs.myof.data['prop_fibrosis']) *
                      self.hs.myof.data['prop_myofilaments'])

        # If it is new beat, calculate cycle metrics
        if ((new_beat > 0) and (self.last_heart_beat_time > 0)):
            # Prune data to last cycle

            # Deduce the fields that are needed
            cycle_fields = ['time', 'pressure_ventricle',
                            'pressure_arteries', 'volume_ventricle']
            if hasattr(self.hs, 'ener'):
                cycle_fields = cycle_fields + ['ener_flux_ATP_consumed']

            d_cycle = self.sim_data[cycle_fields]

            d_cycle = d_cycle[d_cycle['time'].between(
                self.last_heart_beat_time, self.data['time'])]

            self.data['stroke_volume'] = d_cycle['volume_ventricle'].max() - \
                d_cycle['volume_ventricle'].min()

            self.data['ejection_fraction'] = self.data['stroke_volume'] / \
                d_cycle['volume_ventricle'].max()

            self.data['stroke_work'] = self.return_stroke_work(d_cycle)

            self.data['mean_arterial_pressure'] = \
                d_cycle['pressure_arteries'].mean()

            if hasattr(self.hs, 'ener'):
                self.data['stroke_cost'] = \
                    -d_cycle['ener_flux_ATP_consumed'].sum() * time_step * \
                        self.hs.myof.implementation['delta_G_ATP']

                if (self.data['stroke_cost'] > 0):
                    self.data['efficiency'] = self.data['stroke_work'] / \
                        self.data['stroke_cost']
                else:
                    self.data['efficiency'] = np.nan

    def return_flows(self, v, time_step):
        """ return flows between compartments """

        # Calculate pressure in each compartment
        p = np.zeros(self.model['no_of_compartments'])
        for i in np.arange(len(p)-1):
            p[i] = (v[i]-self.data['s'][i]) / self.data['compliance'][i]
        p[-1] = self.return_lv_pressure(v[-1])

        # Add 1 for VAD
        f = np.zeros(self.model['no_of_compartments']+1)
        r = self.data['resistance']
        inert = self.data['inertance']
        last_f = self.data['f']

        for i in np.arange(len(p)):
            f[i] = (1.0 / (r[i] + inert[i]/time_step)) * \
                (p[i-1] - p[i] + (inert[i] * last_f[i] / time_step))

        # Check for VAD
        if (self.va):
            f[-1] = self.va.data['max_flow'] +\
                (p[-1] - p[0]) * self.va.data['pump_slope']

        # Add in the valves
        # Aortic
        if (p[-1] <= p[0]):
            f[0] = (p[-1] - p[0]) * \
                self.data['aortic_insufficiency_conductance']
        # Mitral
        if (p[-1] >= p[-2]):
            f[-2] = (p[-2]-p[-1]) * \
                self.data['mitral_insufficiency_conductance']

        return f

    def evolve_volumes(self, time_step, initial_v):
        """ Evolves volumes for a time-step """
        from scipy.integrate import solve_ivp

        def derivs(t, v):
            dv = np.zeros(self.model['no_of_compartments'])
            flows = self.return_flows(v, time_step)
            for i in np.arange(self.model['no_of_compartments']):
                if (i == (self.model['no_of_compartments']-1)):
                    dv[i] = flows[i] - flows[0] + flows[-1]
                else:
                    dv[i] = flows[i] - flows[i+1]
            return dv

        sol = solve_ivp(derivs, [0, time_step], initial_v)

        # Tidy up negative values
        y = sol.y[:, -1]
        y[np.nonzero(y < 0)] = 0
        # Rest goes in veins
        y[-2] = y[-2] + (self.data['blood_volume'] - np.sum(y))
        return y


#
#    Copyright (C) 2017-2018 European Synchrotron Radiation Facility, Grenoble, France
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#  .
#  The above copyright notice and this permission notice shall be included in
#  all copies or substantial portions of the Software.
#  .
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
#  THE SOFTWARE.

"""CSR rebinning engine implemented in pure python (with bits of scipy !)
"""

__author__ = "Jerome Kieffer"
__contact__ = "Jerome.Kieffer@ESRF.eu"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
__date__ = "19/01/2021"
__status__ = "development"

import logging
logger = logging.getLogger(__name__)
import numpy
from scipy.sparse import csr_matrix
from .preproc import preproc as preproc_np
try:
    from ..ext.preproc import preproc as preproc_cy
except ImportError as err:
    logger.warning("ImportError pyFAI.ext.preproc %s", err)
    preproc = preproc_np
else:
    preproc = preproc_cy

from ..containers import Integrate1dtpl, Integrate2dtpl

class CSRIntegrator(object):

    def __init__(self,
                 image_size,
                 lut=None,
                 empty=0.0):
        """Constructor of the abstract class

        :param size: input image size
        :param lut: tuple of 3 arrays with data, indices and indptr,
                     index of the start of line in the CSR matrix
        :param empty: value for empty pixels
        """
        self.size = image_size
        self.empty = empty
        self.bins = None
        self._csr = None
        self._csr2 = None  # Used for propagating variance
        self.lut_size = 0  # actually nnz
        self.data = None
        self.indices = None
        self.indptr = None
        if lut is not None:
            assert len(lut) == 3
            self.set_matrix(*lut)

    def set_matrix(self, data, indices, indptr):
        """Actually set the CSR sparse matrix content

        :param data: the non zero values NZV
        :param indices: the column number of the NZV
        :param indptr: the index of the start of line"""
        self.data = data
        self.indices = indices
        self.indptr = indptr
        self.lut_size = len(indices)
        self.bins = len(indptr) - 1
        self._csr = csr_matrix((data, indices, indptr), shape=(self.bins, self.size))
        self._csr2 = csr_matrix((data * data, indices, indptr), shape=(self.bins, self.size))  # contains the coef squared, used for variance propagation

    def integrate(self,
                  signal,
                  variance=None,
                  dummy=None,
                  delta_dummy=None,
                  dark=None,
                  flat=None,
                  solidangle=None,
                  polarization=None,
                  absorption=None,
                  normalization_factor=1.0,
                  ):
        """Actually perform the CSR matrix multiplication after preprocessing.

        :param signal: array of the right size with the signal in it.
        :param variance: Variance associated with the signal
        :param dummy: values which have to be discarded (dynamic mask)
        :param delta_dummy: precision for dummy values
        :param dark: noise to be subtracted from signal
        :param flat: flat-field normalization array
        :param flat: solidangle normalization array
        :param polarization: :solidangle normalization array
        :param absorption: :absorption normalization array
        :param normalization_factor: scale all normalization with this scalar
        :return: the preprocessed data integrated as array nbins x 4 which contains:
                    regrouped signal, variance, normalization and pixel count

        Nota: all normalizations are grouped in the preprocessing step.
        """
        shape = signal.shape
        prep = preproc(signal,
                       dark=dark,
                       flat=flat,
                       solidangle=solidangle,
                       polarization=polarization,
                       absorption=absorption,
                       mask=None,
                       dummy=dummy,
                       delta_dummy=delta_dummy,
                       normalization_factor=normalization_factor,
                       empty=self.empty,
                       split_result=4,
                       variance=variance,
                       dtype=numpy.float32)
        prep.shape = numpy.prod(shape), -1
        # logger.warning("prep.shape %s lut_size %s, image_size %s, bins %s", prep.shape, self.lut_size, self.size, self.bins)
        res = numpy.empty((numpy.prod(self.bins), 4), dtype=numpy.float32)
        # logger.warning(self._csr.shape)
        res[:, 0] = self._csr.dot(prep[:, 0])
        if variance is not None:
            res[:, 1] = self._csr2.dot(prep[:, 1])
        res[:, 2] = self._csr.dot(prep[:, 2])
        res[:, 3] = self._csr.dot(prep[:, 3])
        return res

class CsrIntegrator1d(CSRIntegrator):

    def __init__(self,
                 image_size,
                 lut=None,
                 empty=0.0,
                 unit=None,
                 bin_centers=None,
                 ):
        """Constructor of the abstract class for 1D integration

        :param image_size: size of the image
        :param lut: (data, indices, indptr) of the CSR matrix
        :param empty: value for empty pixels
        :param unit: the kind of radial units
        :param bin_center: position of the bin center

        Nota: bins are deduced from bin_centers

        TODO:
        ~/workspace-400/pyFAI/build/lib.linux-x86_64-3.7/pyFAI/azimuthalIntegrator.py in sigma_clip_ng(self, data, npt, correctSolidAngle, polarization_factor, variance, error_model, dark, flat, method, unit, thres, max_iter, dummy, delta_dummy, mask, normalization_factor, metadata, safe, **kwargs)
   3508                         elif (mask is None) and (integr.check_mask):
   3509                             reset = "no mask but CSR has mask"
-> 3510                         elif (mask is not None) and (integr.mask_checksum != mask_crc):
   3511                             reset = "mask changed"
   3512 #                         if (radial_range is None) and (integr.pos0Range is not None):

AttributeError: 'CsrIntegrator1d' object has no attribute 'mask_checksum'

        """
        self.bin_centers = bin_centers
        CSRIntegrator.__init__(self, image_size, lut, empty)
        self.pos0_range = self.pos1_range = self._geometry = None
        self.unit = unit

    def set_geometry(self, geometry):
        from pyFAI.geometry import Geometry
        assert numpy.prod(geometry.detector.shape) == self.size
        assert isinstance(geometry, Geometry)
        self._geometry = geometry

    def set_matrix(self, data, indices, indptr):
        """Actually set the CSR sparse matrix content

        :param data: the non zero values NZV
        :param indices: the column number of the NZV
        :param indptr: the index of the start of line"""

        CSRIntegrator.set_matrix(self, data, indices, indptr)
        assert len(self.bin_centers) == self.bins

    def integrate(self,
                  signal,
                  variance=None,
                  dummy=None,
                  delta_dummy=None,
                  dark=None,
                  flat=None,
                  solidangle=None,
                  polarization=None,
                  absorption=None,
                  normalization_factor=1.0,
                  ):
        """Actually perform the 1D integration

        :param signal: array of the right size with the signal in it.
        :param variance: Variance associated with the signal
        :param dummy: values which have to be discarded (dynamic mask)
        :param delta_dummy: precision for dummy values
        :param dark: noise to be subtracted from signal
        :param flat: flat-field normalization array
        :param flat: solidangle normalization array
        :param polarization: :solidangle normalization array
        :param absorption: :absorption normalization array
        :param normalization_factor: scale all normalization with this scalar
        :return: Integrate1dResult or Integrate1dWithErrorResult object depending on variance

        """
        if variance is None:
            do_variance = False
        else:
            do_variance = True
        trans = CSRIntegrator.integrate(self, signal, variance, dummy, delta_dummy,
                                        dark, flat, solidangle, polarization,
                                        absorption, normalization_factor)
        signal = trans[:, 0]
        variance = trans[:, 1]
        normalization = trans[:, 2]
        count = trans[..., -1]  # should be 3
        mask = (normalization == 0)
        with numpy.errstate(divide='ignore', invalid='ignore'):
            intensity = signal / normalization
            intensity[mask] = self.empty
            if do_variance:
                error = numpy.sqrt(variance) / normalization
                error[mask] = self.empty
            else:
                variance = error = None
        return Integrate1dtpl(self.bin_centers,
                              intensity, error,
                              signal, variance, normalization, count)

    integrate_ng = integrate

    def sigma_clip(self, data, dark=None, dummy=None, delta_dummy=None,
                   variance=None, dark_variance=None,
                   flat=None, solidangle=None, polarization=None, absorption=None,
                   safe=True, error_model=None,
                   normalization_factor=1.0,
                   cutoff=4.0, cycle=5):
        """
        Perform a sigma-clipping iterative filter within each along each row.
        see the doc of scipy.stats.sigmaclip for more descriptions.

        If the error model is "azimuthal": the variance is the variance within a bin,
        which is refined at each iteration, can be costly !

        Else, the error is propagated according to:

        .. math::

            signal = (raw - dark)
            variance = variance + dark_variance
            normalization  = normalization_factor*(flat * solidangle * polarization * absortoption)
            count = number of pixel contributing

        Integration is performed using the CSR representation of the look-up table on all
        arrays: signal, variance, normalization and count

        :param dark: array of same shape as data for pre-processing
        :param dummy: value for invalid data
        :param delta_dummy: precesion for dummy assessement
        :param variance: array of same shape as data for pre-processing
        :param dark_variance: array of same shape as data for pre-processing
        :param flat: array of same shape as data for pre-processing
        :param solidangle: array of same shape as data for pre-processing
        :param polarization: array of same shape as data for pre-processing
        :param safe: if True (default) compares arrays on GPU according to their checksum, unless, use the buffer location is used
        :param normalization_factor: divide raw signal by this value
        :param cutoff: discard all points with |value - avg| > cutoff * sigma. 3-4 is quite common
        :param cycle: perform at maximum this number of cycles. 5 is common.
        :return: namedtuple with "position intensity error signal variance normalization count"
        """
        shape = data.shape
        error_model = error_model.lower() if error_model else ""

        if self._geometry is None:
            raise RuntimeError("Set geometry first")

        prep = preproc(data,
                       dark=dark,
                       flat=flat,
                       solidangle=solidangle,
                       polarization=polarization,
                       absorption=absorption,
                       mask=None,
                       dummy=dummy,
                       delta_dummy=delta_dummy,
                       normalization_factor=normalization_factor,
                       empty=self.empty,
                       split_result=4,
                       variance=variance,
                       dtype=numpy.float32,
                       poissonian=error_model.startswith("pois"))
        prep_flat = prep.reshape((numpy.prod(shape), 4))
        res = self._csr.dot(prep_flat)
        print(cycle)
        for _ in range(cycle):
            msk = res[:, 2] == 0
            avg = res[:, 0] / res[:, 2]
            std = numpy.sqrt(res[:, 1] / res[:, 2])
            avg[msk] = 0
            std[msk] = 0

            avg2d = self._geometry.calcfrom1d(self.bin_centers, avg, shape=shape,
                    dim1_unit=self.unit, correctSolidAngle=False, dummy=0.0)
            std2d = self._geometry.calcfrom1d(self.bin_centers, std, shape=shape,
                    dim1_unit=self.unit, correctSolidAngle=False, dummy=0.0)
            cnt = abs(prep[..., 0] / prep[..., 2] - avg2d) / std2d
            msk2d = numpy.logical_and(numpy.logical_not(numpy.isfinite(cnt)), cnt > cutoff)
            prep[msk2d,:] = 0
            res = self._csr.dot(prep_flat)
        msk = res[:, 2] == 0
        avg = res[:, 0] / res[:, 2]
        std = numpy.sqrt(res[:, 1] / res[:, 2])
        avg[msk] = 0
        std[msk] = 0

        return Integrate1dtpl(self.bin_centers, avg, std, res[:, 0], res[:, 1], res[:, 2], res[:, 3])

class CsrIntegrator2d(CSRIntegrator):

    def __init__(self,
                 image_size,
                 lut=None,
                 empty=0.0,
                 bin_centers0=None,
                 bin_centers1=None):
        """Constructor of the abstract class for 2D integration

        :param size: input image size
        :param lut: tuple of 3 arrays with data, indices and indptr,
                     index of the start of line in the CSR matrix
        :param empty: value for empty pixels
        :param bin_center: position of the bin center

        Nota: bins are deduced from bin_centers0, bin_centers1

        """
        self.bin_centers0 = bin_centers0
        self.bin_centers1 = bin_centers1
        CSRIntegrator.__init__(self, image_size, lut, empty)

    def set_matrix(self, data, indices, indptr):
        """Actually set the CSR sparse matrix content

        :param data: the non zero values NZV
        :param indices: the column number of the NZV
        :param indptr: the index of the start of line"""

        CSRIntegrator.set_matrix(self, data, indices, indptr)
        assert len(self.bin_centers0) * len(self.bin_centers1) == len(indptr) - 1
        self.bins = (len(self.bin_centers0), len(self.bin_centers1))

    def integrate(self,
                  signal,
                  variance=None,
                  dummy=None,
                  delta_dummy=None,
                  dark=None,
                  flat=None,
                  solidangle=None,
                  polarization=None,
                  absorption=None,
                  normalization_factor=1.0):
        """Actually perform the 2D integration

        :param signal: array of the right size with the signal in it.
        :param variance: Variance associated with the signal
        :param dummy: values which have to be discarded (dynamic mask)
        :param delta_dummy: precision for dummy values
        :param dark: noise to be subtracted from signal
        :param flat: flat-field normalization array
        :param flat: solidangle normalization array
        :param polarization: :solidangle normalization array
        :param absorption: :absorption normalization array
        :param normalization_factor: scale all normalization with this scalar
        :return: Integrate2dtpl namedtuple: "radial azimuthal intensity error signal variance normalization count"

        """
        if variance is None:
            do_variance = False
        else:
            do_variance = True
        trans = CSRIntegrator.integrate(self, signal, variance, dummy, delta_dummy,
                                        dark, flat, solidangle, polarization,
                                        absorption, normalization_factor)
        trans.shape = self.bins + (-1,)

        signal = trans[..., 0]
        variance = trans[..., 1]
        normalization = trans[..., 2]
        count = trans[..., -1]  # should be 3
        mask = (normalization == 0)
        with numpy.errstate(divide='ignore', invalid='ignore'):
            intensity = signal / normalization
            intensity[mask] = self.empty
            if do_variance:
                error = numpy.sqrt(variance) / normalization
                error[mask] = self.empty
            else:
                variance = error = None
        return Integrate2dtpl(self.bin_centers0, self.bin_centers1,
                              intensity, error,
                              signal, variance, normalization, count)


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

"""
    ***
    UVW
    ***
"""

__author__ = 'Alan Loh'
__copyright__ = 'Copyright 2020, nenupy'
__credits__ = ['Alan Loh']
__maintainer__ = 'Alan'
__email__ = 'alan.loh@obspm.fr'
__status__ = 'Production'
__all__ = [
    'UVW'
]

import numpy as np
from astropy.coordinates import SkyCoord
import astropy.units as u
from astropy.time import Time
from pyproj import Transformer

from nenupy.instru import ma_info
from nenupy.astro import (
    lst,
    lha,
    toFK5,
    eq_zenith,
    wavelength
)

import logging
log = logging.getLogger(__name__)

# ============================================================= #
# ---------------------------- UVW ---------------------------- #
# ============================================================= #
class UVW(object):
    """
    """

    def __init__(self, times, mas, freqs=None):
        self.bsl_xyz = None
        self.times = times
        self.freqs = freqs
        self.mas = mas

    # --------------------------------------------------------- #
    # --------------------- Getter/Setter --------------------- #
    @property
    def mas(self):
        return self._mas
    @mas.setter
    def mas(self, m):
        if not isinstance(m, np.ndarray):
            raise TypeError(
                'mas should be a numpy array.'
            )
        ma_pos = np.array([a.tolist() for a in ma_info['pos']])
        available_mas = np.arange(ma_pos.shape[0])
        antpos = ma_pos[np.isin(available_mas, m)]
        # RGF93 to ITRF97
        # See http://epsg.io/?q=itrf97 to find correct EPSG
        t = Transformer.from_crs(
            crs_from='EPSG:2154', # RGF93
            crs_to='EPSG:4896'# ITRF2005
        )
        antpos[:, 0], antpos[:, 1], antpos[:, 2] = t.transform(
            xx=antpos[:, 0],
            yy=antpos[:, 1],
            zz=antpos[:, 2]
        )

        xyz = antpos[..., None]
        xyz = xyz[:, :, 0][:, None]
        # xyz = xyz - xyz.transpose(1, 0, 2)
        xyz = xyz.transpose(1, 0, 2) - xyz
        # self.bsl = xyz[np.triu_indices(m.size)]
        self.bsl = xyz[np.tril_indices(m.size)]
        self._mas = m
        return

    @property
    def times(self):
        """ Times at which the UVW must be computed.

            :setter: Times

            :getter: Times

            :type: :class:`~astropy.time.Time`
        """
        return self._times
    @times.setter
    def times(self, t):
        if not isinstance(t, Time):
            raise TypeError(
                'times must be a Time instance.'
            )
        self._times = t
        return

    @property
    def freqs(self):
        return self._freqs
    @freqs.setter
    def freqs(self, f):
        if f is None:
            self._freqs = None
        else:
            if not isinstance(f, u.Quantity):
                f *= u.MHz
            if f.isscalar:
                f = np.array([f.value]) * u.MHz
            self._freqs = f
        return

    @property
    def uvw(self):
        """ UVW in meters.

            :getter: (times, baselines, UVW)

            :type: :class:`~numpy.ndarray`
        """
        if not hasattr(self, '_uvw'):
            raise Exception(
                'Run .compute() first.'
            )
        return self._uvw

    @property
    def uvw_wave(self):
        """ UVW in lambdas.

            :getter: (times, freqs, baselines, UVW)

            :type: :class:`~numpy.ndarray`
        """
        if not hasattr(self, '_uvw'):
            raise Exception(
                'Run .compute() first.'
            )
        if self.freqs is None:
            raise ValueError(
                'No frequency input, fill self.freqs.'
            )
        lamb = wavelength(self.freqs).value
        na = np.newaxis
        return self._uvw[:, na, :, :]/lamb[na, :, na, na]

    # --------------------------------------------------------- #
    # ------------------------ Methods ------------------------ #
    def compute(self, phase_center=None):
        r""" Compute the UVW at a given ``phase_center`` for all
            the :attr:`~nenupy.crosslet.uvw.UVW.times` and baselines
            formed by :attr:`~nenupy.crosslet.uvw.UVW.mas`.

            :param phase_center: Observation phase center. If
                ``None``, local zenith is considered as phase
                center for all :attr:`~nenupy.crosslet.uvw.UVW.times`.
            :type phase_center: :class:`~astropy.coordinates.SkyCoord`

            UVW are computed such as:

            .. math::
                \pmatrix{
                    u \\
                    v \\
                    w
                } =
                \pmatrix{
                    \sin(h) & \cos(h) & 0\\
                    -\sin(\delta) \cos(h) & \sin(\delta) \sin(h) & \cos(\delta)\\
                    \cos(\delta)\cos(h) & -\cos(\delta) \sin(h) & \sin(\delta)
                }
                \pmatrix{
                    \Delta x\\
                    \Delta y\\
                    \Delta z
                }

            :math:`u`, :math:`v`, :math:`w` are in meters. :math:`h`
            is the hour angle (see :func:`~nenupy.astro.astro.lha`)
            at which the phase center is observed, :math:`\delta`
            is the phase center's declination, :math:`(\Delta x,
            \Delta y, \Delta z)` are the baselines projections
            with the convention of :math:`x` to the South, :math:`y`
            to the East and :math:`z` to :math:`\delta = 90` deg.

            Result of the computation are stored as a :class:`~numpy.ndarray`
            in :attr:`~nenupy.crosslet.uvw.UVW.uvw` whose shape is
            (times, cross-correlations, 3), 3 being :math:`(u, v, w)`.
            """
        # Phase center
        if phase_center is None:
            log.info(
                'UVW phase centered at local zenith.'
            )
            phase_center = eq_zenith(self.times)
        else:
            if not isinstance(phase_center, SkyCoord):
                raise TypeError(
                    'phase_center should be a SkyCoord object'
                )
            if phase_center.isscalar:
                ones = np.ones(self.times.size)
                ra_tab = ones * phase_center.ra
                dec_tab = ones * phase_center.dec
                phase_center = SkyCoord(ra_tab, dec_tab)
            else:
                if phase_center.size != self.times.size:
                    raise ValueError(
                        'Size of phase_center != times'
                    )
            log.info(
                'UVW phase centered at RA={}, Dec={}'.format(
                    phase_center.ra[0].deg,
                    phase_center.dec[0].deg
                )
            )
        # Hour angles
        lstTime = lst(
            time=self.times,
            kind='apparent'
        )
        phase_center = toFK5(
            skycoord=phase_center,
            time=self.times
        )
        ha = lha(
            lst=lstTime,
            skycoord=phase_center
        )

        # Transformations
        self._uvw = np.zeros(
            (
                self.times.size,
                self.bsl.shape[0],
                3
            )
        )
        xyz = np.array(self.bsl).T
        # rot = np.radians(-90) # x to the south, y to the east
        # rotation = np.array(
        #     [
        #         [ np.cos(rot), np.sin(rot), 0],
        #         [-np.sin(rot), np.cos(rot), 0],
        #         [ 0,           0,           1]
        #     ]
        # )
        for i in range(self.times.size):
            sr = np.sin(ha[i].rad)
            cr = np.cos(ha[i].rad)
            sd = np.sin(phase_center.dec[i].rad)
            cd = np.cos(phase_center.dec[i].rad)
            rot_uvw = np.array([
                [    sr,     cr,  0],
                [-sd*cr,  sd*sr, cd],
                [ cd*cr, -cd*sr, sd]
            ])
            # self.uvw[i, ...] = - np.dot(
            #     np.dot(rot_uvw, xyz).T,
            #     rotation
            # )
            self._uvw[i, ...] = - np.dot(rot_uvw, xyz).T
        return

    @classmethod
    def from_tvdata(cls, tvdata):
        """
        """
        from nenupy.crosslet import TV_Data
        if not isinstance(tvdata, TV_Data):
            raise TypeError(
                'tvdata must be a TV_Data instance'
            )
        uvw = cls(
            times=tvdata.times,
            freqs=tvdata.freqs,
            mas=tvdata.mas
        )
        uvw.compute()
        return uvw

    @classmethod
    def from_xstdata(cls, xstdata):
        """
        """
        from nenupy.crosslet import XST_Data
        if isinstance(xstdata, XST_Data):
            pass
        elif isinstance(xstdata, str):
            xstdata = XST_Data(xstdata)
        else:
            raise TypeError(
                'xstdata must be a XST_Data instance or XST file.'
            )
        uvw = cls(
            times=xstdata.times,
            freqs=xstdata.freqs,
            mas=xstdata.mas
        )
        uvw.compute()
        return uvw

    # def rephase(self, ra, dec):
    #     """
    #     """
    #     self._uvw
    # --------------------------------------------------------- #
    # ----------------------- Internal ------------------------ #

# ============================================================= #

import numpy,import astropy,from astropy


from __future__ import absolute_import
from ._hydrodiskstars import c_hydrodiskstars
from . import data

class hydrodiskstars:

	r"""
	A stellar migration scheme informed by the ``h277`` simulation, a zoom-in
	hydrodynamic simulation of a Milky Way like galaxy ran from cosmological
	initial conditions (a part of the ``g14`` simulation suite, Christensen et
	al 2012 [1]_).

	**Signature**: vice.toolkit.hydrodisk.hydrodiskstars(radial_bins, N = 1e5,
	mode = "diffusion")

	.. versionadded:: 1.2.0

	.. warning:: Simulations which adopt this model that run for longer than
		13.2 Gyr are not supported. Stellar populations in the built-in
		hydrodynamical simulation data span 13.2 Gyr of ages; simulations on
		longer timescales are highly likely to produce a
		``segmentation fault``.

	Parameters
	----------
	radial_bins : array-like [elements must be positive real numbers]
		The bins in galactocentric radius in kpc describing the disk model.
		This must extend from 0 to at least 20. Need not be sorted in any
		way. Will be stored as an attribute.
	N : int [default : 1e5]
		An approximate number of star particles from the hydrodynamical
		simulation to include in the sample of candidate analogs. Their data
		are not stored in a single file, but split across random subsamples to
		decrease computational overhead when the full sample is not required.
		In practice, this number should be slightly larger than the number of
		(relevant) stellar populations simulated by a multizone model.

		.. note:: There are 3,102,519 star particles available for this
			object. Any more stellar populations than this would oversample
			these data.

	mode : str [case-insensitive] or ``None`` [default : "diffusion"]
		The attribute 'mode', initialized via keyword argument.

	Attributes
	----------
	radial_bins : list
		The bins in galactocentric radius in kpc describing the disk model.
	analog_data : dataframe
		The raw star particle data from the hydrodynamical simulation.
	analog_index : int
		The index of the star particle acting as the current analog. -1 if the
		analog has not yet been set (see note below under `Calling`_).
	mode : str or ``None``
		The mode of stellar migration, describing the approximation of how
		stars move from birth to final radii. Either "diffusion", "sudden", or
		"linear". See property docstring for more details.

		.. note:: Only subclasses may set this attribute ``None``, in which
			case it is assumed that a custom migration approximation is
			employed by an overridden ``__call__`` function. In this case,
			if this attribute is not set to ``None``, multizone simulations
			will *still* use the approximation denoted by this property.

	Calling
	-------
	As all stellar migration prescriptions must, this object can be called
	with three parameters, in the following order:

		zone : int
			The zone index of formation of the stellar population. Must be
			non-negative.
		tform : float
			The time of formation of the stellar population in Gyr.
		time : float
			The simulation time in Gyr (i.e. not the age of the star particle).

	.. note:: The search for analog star particles is ran when the formation
		time and simulation time are equal. Therefore, calling this object
		with the second and third parameters equal resets the star particle
		acting as the analog, and the data for the corresponding star particle
		can then be accessed via the attribute ``analog_index``.

	Functions
	---------
	decomp_filter : [instancemethod]
		Filter the star particles based on their kinematic decomposition.

	Raises
	------
	* ValueError
		- Minimum radius does not equal zero
		- Maximum radius < 20
	* ScienceWarning
		- This object is called with a time larger than 13.2 Gyr
		- The number of analog star particles requested is larger than the
		  number available from the hydrodynamical simulation (3,102,519)

	Notes
	-----
	This object requires VICE's supplementary data sample, available in its
	GitHub repository at ./vice/toolkit/hydrodisk/data.  The first time a
	``hydrodiskstars`` object is constructed, VICE will download
	the additional data automatically. If this process fails, it may be due to
	not having administrator's privileges on your system; in this event, users
	should speak with their administrator, who would then be able to download
	their data by running the following on their system:

	>>> import vice
	>>> vice.toolkit.hydrodisk.data.download()

	This migration scheme works by assigning each stellar population in the
	simulation an analog star particle from the hydrodynamical simulation. The
	analog is randomly drawn from a sample of star particles which formed at
	a similar radius and time, and the stellar population then assumes the
	change in orbital radius of its analog.

	VICE first searches for analogs in the ``h277`` data for star particles
	which formed at a radius of :math:`R \pm` 250 pc and at a time of
	:math:`T \pm` 250 Myr. If no analogs are found that satisfy this
	requirement, the search is widened to :math:`R \pm` 500 pc and
	:math:`T \pm` 500 Myr. If still no analog is found, then the time
	restriction of :math:`T \pm` 500 Myr is maintained, and VICE finds the star
	particle with the smallest difference in birth radius, assigning it as the
	analog. These values parameterizing this search algorithm are declared in
	``vice/src/toolkit/hydrodiskstars.h`` in the VICE source tree. For further
	details, see "Milky Way-Like Galaxies" under VICE's science documentation.

	This object can be subclassed to implement a customized migration
	approximation by overriding the ``__call__`` function. However, in this
	case, users must also set the attribute ``mode`` to ``None``. If this
	requirement is not satisfied, multizone simulations will **still** use the
	approximation denoted by the ``mode`` attribute, **not** their overridden
	``__call__`` function.

	The ``h277`` galaxy had a weak and transient bar, but does not have one at
	the present day. This is one notable difference between it and the Milky
	Way.

	Example Code
	------------
	>>> from vice.toolkit.hydrodisk import hydrodiskstars
	>>> import numpy as np
	>>> example = hydrodiskstars(np.linspace(0, 20, 81), N = 5e5)
	>>> example.radial_bins
	[0.0,
	 0.25,
	 0.5,
	 ...
	 19.5,
	 19.75,
	 20.0]
	>>> example.analog_data.keys()
	['id', 'tform', 'rform', 'rfinal', 'zfinal', 'vrad', 'vphi', 'vz']
	>>> example.analog_index
	-1
	>>> example(5, 7.2, 7.2)
	5
	>>> example.analog_index
	200672
	>>> example.analog_data["vrad"][example.analog_index]
	5.6577
	>>> example.mode
	"diffusion"

	.. [1] Christensen et al. (2012), MNRAS, 425, 3058
	"""

	def __init__(self, rad_bins, N = 1e5, mode = "diffusion"):
		if not data._h277_exists():
			print("VICE supplementary data required, downloading now.")
			print("You will not need to repeat this process.")
			data.download()
		else: pass
		self.__c_version = c_hydrodiskstars(rad_bins, N = N, mode = mode)

	def __call__(self, zone, tform, time):
		return self.__c_version.__call__(zone, tform, time)

	def __enter__(self):
		# Opens a with statement
		return self

	def __exit__(self, exc_type, exc_value, exc_tb):
		# Raises all exceptions inside a with statement
		return exc_value is None

	def __object_address(self):
		r"""
		Returns the memory address of the HYDRODISKSTARS object in C. For
		internal usage only; usage of this function by the user is strongly
		discouraged.
		"""
		return self.__c_version.object_address()

	@property
	def radial_bins(self):
		r"""
		Type : list [elements are positive real numbers]

		The bins in galactocentric radius in kpc describing the disk model.
		Must extend from 0 to at least 20 kpc. Need not be sorted in any way
		when assigned.

		Example Code
		------------
		>>> from vice.toolkit.hydrodisk import hydrodiskstars
		>>> import numpy as np
		>>> example = hydrodiskstars([0, 5, 10, 15, 20])
		>>> example.radial_bins
		[0, 5, 10, 15, 20]
		>>> example.radial_bins = list(range(31))
		>>> example.radial_bins
		[0,
		 1,
		 2,
		 ...
		 17,
		 18,
		 19,
		 20]
		"""
		return self.__c_version.radial_bins

	@radial_bins.setter
	def radial_bins(self, value):
		self.__c_version.radial_bins = value

	@property
	def analog_data(self):
		r"""
		Type : dataframe

		The star particle data from the hydrodynamical simulation. The
		following keys map to the following data:

			- id:      	The IDs of each star particle
			- tform:   	The time the star particle formed in Gyr
			- rform:   	The radius the star particle formed at in kpc
			- rfinal:  	The radius the star particle ended up at in kpc
			- zform: 	The disk midplane distance in kpc at the time of
			  formation.
			- zfinal:  	The disk midplane distance in kpc at the end of the
			  simulation
			- vrad:     The radial velocity of the star particle at the end of
			  the simulation in km/sec
			- vphi:     The azimuthal velocity of the star particle at the end
			  of the simulation in km/sec
			- vz: 		The velocity perpendicular to the disk midplane at the
			  end of the simulation in km/sec
			- decomp: 	An integer denoting which kinematic subclass the star
			  particle belongs to (1: thin disk, 2: thick disk, 3: bulge,
			  4: pseudobulge, 5: halo).

		Example Code
		------------
		>>> from vice.toolkit.hydrodisk import hydrodiskstars
		>>> import numpy as np
		>>> example = hydrodiskstars(np.linspace(0, 20, 81))
		>>> example.analog_data.keys()
		['id', 'tform', 'rform', 'rfinal', 'zfinal', 'vrad', 'vphi', 'vz']
		>>> example.analog_data["rfinal"][:10]
		[2.0804,
		 14.9953,
		 2.2718,
		 15.1236,
		 2.3763,
		 0.9242,
		 9.0908,
		 0.1749,
		 8.415,
		 20.1452]
		"""
		return self.__c_version.analog_data

	@property
	def analog_index(self):
		r"""
		Type : int

		The index of the analog in the hydrodynamical simulation star particle
		data. -1 if it has not yet been assigned.

		.. note:: Calling this object at a given zone with the formation time
			and the simulation time equal resets the star particle acting as
			the analog.

		Example Code
		------------
		>>> from vice.toolkit.hydrodisk import hydrodiskstars
		>>> import numpy as np
		>>> example = hydrodiskstars(np.linspace(0, 20, 81))
		>>> example.analog_index
		-1 # no analog yet
		>>> example(2, 1, 1) # final two arguments equal resets analog
		15745
		>>> example(10, 4, 4)
		10
		>>> example.analog_index
		101206
		>>> example.analog_data["rfinal"][example.analog_index]
		2.6411
		>>> example.analog_data["vrad"][example.analog_data]
		92.2085
		"""
		return self.__c_version.analog_index

	@property
	def mode(self):
		r"""
		Type : str [case-insensitive] or ``None``

		Default : "diffusion"

		Recognized Values
		-----------------
		The following is a breakdown of how stellar populations migrate in
		multizone simulations under each approximation.

		- "diffusion"
			The orbital radius at times between birth and 13.2 Gyr are assigned
			via a sqrt(time) dependence, the mean displacement if stars
			migrated according to a random walk. Stellar populations spiral
			inward or outward with a smooth time dependence in this model.
		- "linear"
			Orbital radii at times between birth and 13.2 Gyr are assigned via
			linear interpolation. Stellar populations therefore spiral
			uniformly inward or outward from birth to final radii, with orbital
			radius changing more slowly for young stars than in the diffusion
			model, but more quickly for old stars.
		- "sudden"
			The time of migration is randomly drawn from a uniform distribution
			between when a stellar population is born and 13.2 Gyr. At times
			prior to this, it is at its radius of birth; at subsequent times,
			it is at its final radius. Stellar populations therefore spend no
			time at intermediate radii.
		- ``None``
			Only supported for subclasses of the hydrodiskstars object, this
			should be used when the user intends to override the built-in
			migration assumptions and implement a different one. In these
			cases, the ``__call__`` method should be overridden in addition to
			this attribute being set to ``None``.

		.. note:: The "post-processing" model from Johnson et al. (2021) [1]_
			corresponds to setting the attribute ``simple = True`` in the
			``milkyway`` object, inherited from the base class ``multizone``.

		Example Code
		------------
		>>> from vice.toolkit.hydrodisk import hydrodiskstars
		>>> import numpy as np
		>>> example = hydrodiskstars(np.linspace(0, 20, 81))
		>>> example.mode
		'diffusion'
		>>> example.mode = "sudden"
		>>> example.mode = "linear"

		.. [1] Johnson et al. (2021), 2103.09838
		"""
		return self.__c_version.mode

	@mode.setter
	def mode(self, value):
		if value is None:
			if type(self) != hydrodiskstars:
				self.__c_version.mode = None
			else:
				raise ValueError("""None-Type value for attribute 'mode' only \
supported for subclasses of hydrodiskstars object.""")
		else:
			self.__c_version.mode = value

	def decomp_filter(self, values):
		r"""
		Filter the star particles from the hydrodynamic simulation based on
		the kinematic decomposition.

		Parameters
		----------
		values : ``list`` [elements of type ``int``]
			The integer values of the "decomp" column in the ``analog_data``
			attribute to base the filter on. Those with a decomposition tag
			equal to one of the values in this list will pass the filter and
			remain in the sample.

			.. note:: The integer values mean that an individual star particle
				has kinematics associated with the following sub-populations:

					- 1: thin disk
					- 2: thick disk
					- 3: bulge
					- 4: pseudobulge
					- 5: halo

		Example Code
		------------
		>>> from vice.toolkit.hydrodisk import hydrodiskstars
		>>> import numpy as np
		>>> example = hydrodiskstars(np.linspace(0, 20, 81))
		>>> len(example.analog_data['id'])
		102857
		>>> example.decomp_filter([1, 2]) # disk stars only
		>>> len(example.analog_data['id'])
		57915
		>>> all([i in [1, 2] for i in example.analog_data['decomp']])
		True
		>>> example = hydrodiskstars(np.linspace(0, 20, 81))
		>>> len(example.analog_data['id'])
		102857
		>>> example.decomp_filter([3, 4]) # bulge stars only
		>>> len(example.analog_data['id'])
		44942
		>>> all([i in [3, 4] for i in example.analog_data['decomp']])
		True
		"""
		self.__c_version.decomp_filter(values)


from __future__ import division
import os
import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.nn.functional import upsample,normalize
from ...nn import PAM_Module
from ...nn import CAM_Module
from .base import BaseNet
from torch.nn.parallel.scatter_gather import gather

__all__ = ['DANet', 'get_danet']
def random_rectangle_mask_gene(img_size, batch_size, num_rect, device, max_size_factor):

    msk = torch.ones((batch_size, 1, img_size[0], img_size[1]), device=device).float()

    for batch_idx in range(batch_size):
        for i in range(num_rect):
            pos = np.random.rand(2)
            size = np.maximum(np.random.rand(2), 0.1/max_size_factor) * max_size_factor

            h_start = int(pos[0] * img_size[0])
            w_start = int(pos[1] * img_size[1])
            h = np.minimum(int(size[0] * img_size[0]), img_size[0]-h_start)
            w = np.minimum(int(size[1] * img_size[1]), img_size[1]-w_start)

            msk[batch_idx, :, h_start:h_start+h, w_start:w_start+w] = 0.

    return msk

#maxpooling as inpainting module
def MaxpoolingAsInpainting(x, x_fore_augment):
    #print(x.detach().cpu().numpy().shape)
    #print(str(4))
    x_fore_augments = gather(x_fore_augment, 0, dim=0)
    pred = torch.max(x_fore_augment[0], 1)[1].unsqueeze(1)

    #print(str(5))
    #print(x_fore_augment[0].detach().cpu().numpy().shape)

    x_fore_pool = F.interpolate(pred.float(), size=x.size()[2:4]).detach().clone()
    fore_msk = (x_fore_pool > 0.5).float().detach().clone()
    fore_msk = F.max_pool2d((fore_msk).detach().clone(), kernel_size=3, stride=1, padding=1)
    bkgd_msk_prev = (1.-fore_msk).detach().clone()
    bkgd_msk =  bkgd_msk_prev.detach().clone()
    x_new = x.detach().clone() * bkgd_msk.detach()[0].clone()
    x_patch = x.detach().clone() * 0.

    while torch.mean(1.-bkgd_msk).item() > 0.0001 and torch.min(torch.mean(bkgd_msk, dim=(1, 2, 3))) > 0.:
        x_new =0.5* F.max_pool2d((x_new).detach().clone(), kernel_size=3, stride=1, padding=1)+ 0.5*F.avg_pool2d((x_new).detach().clone(), kernel_size=3, stride=1, padding=1)
        bkgd_msk = F.max_pool2d((bkgd_msk_prev).detach().clone(),kernel_size=3, stride=1, padding=1)
        x_patch = (x_patch + (bkgd_msk-bkgd_msk_prev)*x_new).detach().clone()
        bkgd_msk_prev = bkgd_msk.detach().clone()

    return x_patch,pred

class PSPModule(nn.Module):
    def __init__(self, features, out_features=1024, sizes=(1, 2, 3, 6)):
        super().__init__()
        self.stages = []
        self.stages = nn.ModuleList([self._make_stage(features, size) for size in sizes])
        self.bottleneck = nn.Conv2d(features * (len(sizes) + 1), out_features, kernel_size=1)
        self.relu = nn.ReLU()

    def _make_stage(self, features, size):
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        conv = nn.Conv2d(features, features, kernel_size=1, bias=False)
        return nn.Sequential(prior, conv)

    def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        priors = [F.interpolate(input=stage(feats), size=(h, w), mode='bilinear', align_corners=False) for stage in self.stages] + [feats]
        bottle = self.bottleneck(torch.cat(priors, 1))
        return self.relu(bottle)

class PSPUpsample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.PReLU()
        )

    def forward(self, x):
        h, w = 2 * x.size(2), 2 * x.size(3)
        p = F.interpolate(input=x, size=(h, w), mode='bilinear', align_corners=False)
        return self.conv(p)

class PSPNet(nn.Module):
    def __init__(self, n_classes, sizes, psp_size, backend,pretrained):
        super().__init__()
        self.feats = backend
        self.psp = PSPModule(psp_size, 1024, sizes)
        self.drop_1 = nn.Dropout2d(p=0.3)

        self.up_1 = PSPUpsample(1024, 256)
        self.up_2 = PSPUpsample(256, 64)
        self.up_3 = PSPUpsample(64, 64)

        self.drop_2 = nn.Dropout2d(p=0.15)
        self.final = nn.Sequential(
            nn.Conv2d(64, n_classes, kernel_size=1),
        )

    def forward(self, x, training, device, gt_fore, use_gt_fore):

        f = self.feats
        p = self.psp(f)
        p = self.drop_1(p)

        p = self.up_1(p)
        p = self.drop_2(p)

        p = self.up_2(p)
        p = self.drop_2(p)

        p = self.up_3(p)
        p = self.drop_2(p)

        return self.final(p)

class Upsample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.PReLU()
        )

    def forward(self, x):
        h, w = 2 * x.size(2), 2 * x.size(3)
        p = F.interpolate(input=x, size=(h, w), mode='bilinear', align_corners=False)
        return self.conv(p)

def load_weights_sequential(target, source_state):
    model_to_load= {k: v for k, v in source_state.items() if k in target.state_dict().keys()}
    target.load_state_dict(model_to_load, strict=False)

def conv3x3(in_planes, out_planes, stride=1, dilation=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, dilation=dilation, bias=False)

class DANet(BaseNet):
    r"""Fully Convolutional Networks for Semantic Segmentation

    Parameters
    ----------
    nclass : int
        Number of categories for the training dataset.
    backbone : string
        Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50',
        'resnet101' or 'resnet152').
    norm_layer : object
        Normalization layer used in backbone network (default: :class:`mxnet.gluon.nn.BatchNorm`;

    Reference:

        Long, Jonathan, Evan Shelhamer, and Trevor Darrell. "Fully convolutional networks
        for semantic segmentation." *CVPR*, 2015

    """
    def __init__(self, nclass, backbone, aux=False, se_loss=False, norm_layer=nn.BatchNorm2d, **kwargs):
        super(DANet, self).__init__(nclass, backbone, aux, se_loss, norm_layer=norm_layer, **kwargs)
        #print(str(10))
        inter_channels=512
        self.head = DANetHead(2048, nclass, norm_layer)
        self.conv511 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
                                   norm_layer(inter_channels),
                                   nn.ReLU())

        #self.conv6 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, out_channels, 1))
        #self.conv10 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(512, 2, 1))

        self.convforeout = nn.Sequential(
            #nn.Conv2d(1024, 512, 3, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            Upsample(512, 256),
            Upsample(256, 128),
            nn.Conv2d(128, 1, 1, bias=False)
        )

        self.convforeout2 = nn.Sequential(
            #nn.Conv2d(64, 512, 3, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            Upsample(512, 256),
            Upsample(256, 128),
            nn.Conv2d(128, 1, 1, bias=False)
        )
        self.convforeout3 = nn.Sequential(
            #nn.Conv2d(64, 512, 3, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, 3, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            Upsample(512, 256),
            Upsample(256, 128),
            nn.Conv2d(128, 1, 1, bias=False)
        )

    def forward(self, x):
        #print(str(0))
        imsize = x.size()[2:]
        _, _, c3, c4 = self.base_forward(x)
        #print(str(100))
        x = self.head(c4)
        x = list(x)

        #x[0] = upsample(x[0], imsize, **self._up_kwargs)
        x[0]=self.conv511(x[0])
        x[1] = upsample(x[1], imsize, **self._up_kwargs)
        x[2] = upsample(x[2], imsize, **self._up_kwargs)
        x[3] = upsample(x[3], imsize, **self._up_kwargs)
        #2.maxpooling as inpainting
        final_featuremap,x_fore_augment=MaxpoolingAsInpainting(x[0],[x[1],x[2],x[3]])

        x1 = final_featuremap * (F.interpolate(x_fore_augment.float(), size=x[0].size()[2:4]) > 0.5).float() +\
            x[0] * (F.interpolate(x_fore_augment.float(), size=x[0].size()[2:4]) < 0.5).float()

        #=output=self.conv10(x1)

        x_fore_pred = self.convforeout(x1+x[0])
        x_fore_pred = torch.sigmoid(x_fore_pred)
        output = F.interpolate(x_fore_pred, scale_factor=2)

        #output = upsample(output, imsize, **self._up_kwargs)

        outputs = [output]
        outputs.append(x[1])
        outputs.append(x[2])
        outputs.append(x[3])
        outputs.append(F.interpolate(torch.sigmoid(self.convforeout2(x[0])),scale_factor=2))
        outputs.append(F.interpolate(torch.sigmoid(self.convforeout3(x1)),scale_factor=2))
        #print("output length is:"+str(len(outputs))+"\n")
        return tuple(outputs)

class DANetHead(nn.Module):
    def __init__(self, in_channels, out_channels, norm_layer):
        super(DANetHead, self).__init__()
        inter_channels = in_channels // 4
        self.conv5a = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
                                   norm_layer(inter_channels),
                                   nn.ReLU())

        self.conv5c = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
                                   norm_layer(inter_channels),
                                   nn.ReLU())

        self.sa = PAM_Module(inter_channels)

        self.sc = CAM_Module(inter_channels)
        self.conv51 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
                                   norm_layer(inter_channels),
                                   nn.ReLU())
        self.conv52 = nn.Sequential(nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False),
                                   norm_layer(inter_channels),
                                   nn.ReLU())

        self.conv6 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, out_channels, 1))
        self.conv7 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, out_channels, 1))

        self.conv8 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, out_channels, 1))
        self.conv9 = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(inter_channels, 2, 1))

        #self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2)
        #self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4)

    def _make_layer(self, block, planes, blocks, stride=1, dilation=1):
        downsample = None

        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = [block(self.inplanes, planes, stride, downsample)]
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes, dilation=dilation))

        return nn.Sequential(*layers)

    def forward(self, x):
        #print(str(2))
        feat1 = self.conv5a(x)

        #the foreground branch
        sa_feat = self.sa(feat1)
        sa_conv = self.conv51(sa_feat)
        sa_output = self.conv6(sa_conv)

        feat2 = self.conv5c(x)
        sc_feat = self.sc(feat2)
        sc_conv = self.conv52(sc_feat)
        sc_output = self.conv7(sc_conv)

        feat_sum = sa_conv+sc_conv

        sasc_output = self.conv8(feat_sum)
        #print(str(3))
        #the road branch
        #1. intermediate feature map extraction
        sa_feat2=self.sa(feat1)
        #print(str(1))

        #x_3 = self.layer3(x)
        #x = self.layer4(x_3)
        #final=PSPNet(x,, psp_size=1024, pretrained=False, n_classes=3)
        #use pspnet maybe?
        return (sa_feat2, sasc_output, sa_output, sc_output)

def get_danet(dataset='pascal_voc', backbone='resnet50', pretrained=False,
           root='~/.encoding/models', **kwargs):
    r"""DANet model from the paper `"Dual Attention Network for Scene Segmentation"
    <https://arxiv.org/abs/1809.02983.pdf>`
    """
    acronyms = {
        'pascal_voc': 'voc',
        'pascal_aug': 'voc',
        'pcontext': 'pcontext',
        'ade20k': 'ade',
        'cityscapes': 'cityscapes',
    }
    # infer number of classes
    from ...datasets import datasets, VOCSegmentation, VOCAugSegmentation, ADE20KSegmentation
    model = DANet(20, backbone=backbone, root=root, **kwargs)
    if pretrained:
        from .model_store import get_model_file
        model.load_state_dict(torch.load(
            get_model_file('fcn_%s_%s'%(backbone, acronyms[dataset]), root=root)),
            strict=False)
    return model


import numpy
import simtk.unit
import simtk.openmm as mm

print('OpenMM version: ', mm.version.full_version)
from openmmtools.constants import kB

class CustomizableGHMC(mm.CustomIntegrator):
    """Customizable GHMC: proposal can contain mix of deterministic and stochastic steps.
    If the provided `splitting` is deterministic, then we add the stochastic velocity update to the beginning of the splitting.

    The class of allowable splittings is the same as for the LangevinSplittingIntegrator
    """

    def __init__(self,
                 splitting="V R V",
                 temperature=298.0 * simtk.unit.kelvin,
                 collision_rate=91.0 / simtk.unit.picoseconds,
                 timestep=1.0 * simtk.unit.femtoseconds,
                 constraint_tolerance=1e-4,
                 ):
        """
        Parameters
        ----------
        splitting : string
            Sequence of R, V substeps to be executed each timestep.
            Forces are only used in V-step. Handle multiple force groups by appending the force group index
            to V-steps, e.g. "V0" will only use forces from force group 0. "V" will perform a step using all forces.

        temperature : numpy.unit.Quantity compatible with kelvin, default: 298.0*simtk.unit.kelvin
           Temperature of heat bath

        collision_rate : numpy.unit.Quantity compatible with 1/picoseconds, default: 91.0/simtk.unit.picoseconds
           Collision rate

        timestep : numpy.unit.Quantity compatible with femtoseconds, default: 1.0*simtk.unit.femtoseconds
           Integration timestep

        constraint_tolerance : float
            Numerical tolerance for solving constraints
        """

        # Compute constants
        kT = kB * temperature
        gamma = collision_rate
        kinetic_energy = "0.5 * m * v * v"

        # Convert splitting string into a list of all-caps strings
        splitting = splitting.upper().split()

        # Count how many times each step appears, so we know how big each R/V/O substep will be
        n_R = sum([letter == "R" for letter in splitting])
        n_V = sum([letter == "V" for letter in splitting])
        n_O = sum([letter == "O" for letter in splitting])

        # If the splitting is deterministic, add a velocity-randomization step to the beginning
        # of the procedure
        if n_O == 0:
            splitting = ["O"] + splitting

        # Check if the splitting string asks for multi-time-stepping.
        # If so, each force group should be integrated for a total length equal to dt
        if len(set([step for step in splitting if step[0] == "V"])) > 1:
            mts = True
            fgs = set([step[1:] for step in splitting if step[0] == "V"])
            n_Vs = dict()
            for fg in fgs:
                n_Vs[fg] = sum([step[1:] == fg for step in splitting])
        else:
            mts = False

        # Do a couple sanity checks on the splitting string
        # Make sure we contain at least one of R, V, O steps
        assert ("R" in splitting)
        assert ("V" in [s[0] for s in splitting])
        assert ("O" in splitting)

        # Make sure it contains no invalid characters
        assert (set(splitting).issubset(set("RVO").union(set(["V{}".format(i) for i in range(32)]))))

        # If the splitting string contains both "V" and a force-group-specific V0,V1,etc.,
        # then raise an error
        if mts and (n_V > 0):
            raise (ValueError("Splitting string includes an evaluation of all forces and "
                              "evaluation of subsets of forces."))

        # Define substep functions
        def R_step():
            # update positions (and velocities, if there are constraints)
            self.addComputePerDof("x", "x + ((dt / {}) * v)".format(n_R))
            self.addComputePerDof("x1", "x")  # save pre-constraint positions in x1
            self.addConstrainPositions()  # x is now constrained
            self.addComputePerDof("v", "v + ((x - x1) / (dt / {}))".format(n_R))
            self.addConstrainVelocities()

        def V_step(fg):
            """Deterministic velocity update, using only forces from force-group fg.
            Parameters
            ----------
            fg : string
                Force group to use in this substep.
                "" means all forces, "0" means force-group 0, etc.
            """
            # update velocities
            if mts:
                self.addComputePerDof("v", "v + ((dt / {}) * f{} / m)".format(n_Vs[fg], fg))
            else:
                self.addComputePerDof("v", "v + (dt / {}) * f / m".format(n_V))
            self.addConstrainVelocities()

        def O_step():
            # measure heat
            self.addComputeSum("old_ke", kinetic_energy)

            # update velocities
            self.addComputePerDof("v", "(a * v) + (b * sqrt(kT / m) * gaussian)")
            self.addConstrainVelocities()

            # measure heat
            self.addComputeSum("new_ke", kinetic_energy)
            self.addComputeGlobal("heat", "heat + (new_ke - old_ke)")

        def substep_function(step_string):
            if step_string == "R":
                R_step()
            elif step_string[0] == "V":
                V_step(step_string[1:])

        def compute_total_energy(name="e_old"):
            self.addComputeSum("ke", kinetic_energy)
            self.addComputeGlobal(name, "ke + energy")

        # Create a new CustomIntegrator
        super(CustomizableGHMC, self).__init__(timestep)

        # Initialize
        self.addGlobalVariable("kT", kT)

        # Velocity mixing parameter: current velocity component
        h = timestep
        self.addGlobalVariable("a", numpy.exp(-gamma * h))

        # Velocity mixing parameter: random velocity component
        self.addGlobalVariable("b", numpy.sqrt(1 - numpy.exp(- 2 * gamma * h)))

        # Positions before application of position constraints
        self.addPerDofVariable("x1", 0)

        # Set constraint tolerance
        self.setConstraintTolerance(constraint_tolerance)

        # Add bookkeeping variables
        self.addGlobalVariable("ke_old", 0)
        self.addGlobalVariable("ke_new", 0)
        self.addGlobalVariable("heat", 0)
        self.addGlobalVariable("shadow_work", 0)
        self.addGlobalVariable("e_old", 0)
        self.addGlobalVariable("e_new", 0)
        self.addGlobalVariable("acc_ratio", 0)
        self.addGlobalVariable("accept", 0)
        self.addGlobalVariable("naccept", 0)
        self.addGlobalVariable("ntrials", 0)
        self.addPerDofVariable("xold", 0)
        self.addPerDofVariable("vold", 0)

        # Compute energy of current state
        self.addUpdateContextState()
        self.addComputePerDof("xold", "x")
        self.addComputePerDof("vold", "v")
        compute_total_energy("e_old")

        # Reset "heat" to zero
        self.addComputeGlobal("heat", "0")

        # Integrate / generate proposal
        for i, step in enumerate(splitting):
            substep_function(step)

        # Compute M-H ratio in terms of energy change during the deterministic steps
        compute_total_energy("e_new")
        self.addComputeGlobal("shadow_work", "(e_new - e_old) - heat")
        self.addComputeGlobal("acc_ratio", "exp(- shadow_work / kT)")

        # Accept / reject : flip momenta upon rejection
        self.addComputeGlobal("accept", "step(acc_ratio - uniform)")
        self.beginIfBlock("accept != 1")
        self.addComputePerDof("x", "xold")
        self.addComputePerDof("v", "-vold")
        self.endBlock()

        # Update acceptance statistics
        self.addComputeGlobal("naccept", "naccept + accept")
        self.addComputeGlobal("ntrials", "ntrials + 1")