import logging
from numba import njit
import numpy as np
from scipy.stats import skewnorm
from strax import exporter
from .pulse import Pulse
from .. import units
from ..load_resource import DummyMap

export, __all__ = exporter()
logging.basicConfig(handlers=[logging.StreamHandler()])
log = logging.getLogger('wfsim.core')
log.setLevel('WARNING')


@export
class S2(Pulse):
    """
    Given temporal inputs as well as number of electrons
    Random generate photon timing and channel distribution.
    """

    def __init__(self, config):
        super().__init__(config)

        self.phase = 'gas'  # To distinguish singlet/triplet time delay.
        self.luminescence_switch_threshold = 100  # When to use simplified model (NOT IN USE)

    @staticmethod
    def inverse_field_distortion_correction(x, y, z, resource):
        """For 1T the pattern map is a data driven one so we need to reverse engineer field distortion correction
        into the simulated positions
        :param x: 1d array of float
        :param y: 1d array of float
        :param z: 1d array of float
        :param resource: instance of resource class
        returns z: 1d array, postions 2d array 
        """
        positions = np.array([x, y, z]).T
        for i_iter in range(6):  # 6 iterations seems to work
            dr = resource.fdc_3d(positions)
            if i_iter > 0:
                dr = 0.5 * dr + 0.5 * dr_pre  # Average between iter
            dr_pre = dr

            r_obs = np.sqrt(x**2 + y**2) - dr
            x_obs = x * r_obs / (r_obs + dr)
            y_obs = y * r_obs / (r_obs + dr)
            z_obs = - np.sqrt(z**2 + dr**2)
            positions = np.array([x_obs, y_obs, z_obs]).T

        positions = np.array([x_obs, y_obs]).T 
        return z_obs, positions
    
    @staticmethod
    def field_distortion_comsol(x, y, z, resource):
        """Field distortion from the COMSOL simulation for the given electrode configuration:
        :param x: 1d array of float
        :param y: 1d array of float
        :param z: 1d array of float
        :param resource: instance of resource class
        returns z: 1d array, postions 2d array 
        """
        positions = np.array([np.sqrt(x**2 + y**2), z]).T
        theta = np.arctan2(y, x)
        r_obs = resource.fd_comsol(positions, map_name='r_distortion_map')
        x_obs = r_obs * np.cos(theta)
        y_obs = r_obs * np.sin(theta)
        
        positions = np.array([x_obs, y_obs]).T 
        return z, positions

    def __call__(self, instruction):
        if len(instruction.shape) < 1:
            # shape of recarr is a bit strange
            instruction = np.array([instruction])

        _, _, t, x, y, z, n_electron, recoil_type, *rest = [
            np.array(v).reshape(-1) for v in zip(*instruction)]
        
        # Reverse engineering FDC
        if self.config['field_distortion_model'] == 'inverse_fdc':
            z_obs, positions = self.inverse_field_distortion_correction(x, y, z, resource=self.resource)
        # Reverse engineering FDC
        elif self.config['field_distortion_model'] == 'comsol':
            z_obs, positions = self.field_distortion_comsol(x, y, z, resource=self.resource)
        else:
            z_obs, positions = z, np.array([x, y]).T

        n_electron = self.get_electron_yield(n_electron=n_electron,
                                             xy_int= np.array([x, y]).T, # maps are in R_true, so orginal position should be here
                                             z_int=z, # maps are in Z_true, so orginal position should be here
                                             config=self.config,
                                             resource=self.resource)

        sc_gain = self.get_s2_light_yield(positions=positions,
                                          config=self.config,
                                          resource=self.resource)
        


        n_photons_per_xy, n_photons_per_ele, self._electron_timings = self.get_n_photons(t=t,
                                                                                         n_electron=n_electron,
                                                                                         z_int=z,
                                                                                         xy_int= np.array([x, y]).T,
                                                                                         sc_gain=sc_gain, 
                                                                                         config=self.config,
                                                                                         resource=self.resource)

        self._instruction = np.repeat(np.arange(len(n_electron)), n_photons_per_xy)
        self._photon_channels = self.photon_channels(n_electron=n_electron,
                                                     z_obs=z_obs,
                                                     positions=positions,
                                                     _instruction=self._instruction,
                                                     config=self.config,
                                                     resource=self.resource)

        # Second generate photon timing
        self._photon_timings = self.photon_timings(positions=positions, 
                                                   n_photons_per_xy=n_photons_per_xy,
                                                   _electron_timings=self._electron_timings,
                                                   n_photons_per_ele=n_photons_per_ele,
                                                   _photon_channels=self._photon_channels,
                                                   phase=self.phase,
                                                   config=self.config,
                                                   resource=self.resource,)

        # Sorting times according to the channel, as non-explicit sorting
        # is performed later and this breaks timing of individual channels/arrays
        sortind = np.argsort(self._photon_channels)

        self._photon_channels = self._photon_channels[sortind]
        self._photon_timings = self._photon_timings[sortind]

        super().__call__()

    @staticmethod
    def get_avg_drift_velocity(z, xy, config, resource):
        """Calculate s2 drift time mean and spread

        :param z: 1d array of z (floats)
        :param xy: 2d array of xy positions (floats)
        :param config: dict with wfsim config
        :param resource: instance of the resource class

        returns array of floats corresponding to average drift velocities from given point to the gate
        """
        if config['enable_field_dependencies']['drift_speed_map']:
            drift_v_LXe = resource.field_dependencies_map(z, xy, map_name='drift_speed_map')  # mm/µs
            drift_v_LXe *= 1e-4  # cm/ns
            drift_v_LXe *= resource.drift_velocity_scaling
        else:
            drift_v_LXe = config['drift_velocity_liquid']
        return drift_v_LXe

    @staticmethod
    def get_s2_drift_time_params(z_int, xy_int, config, resource):
        """Calculate s2 drift time mean and spread

        :param z_int: 1d array of true z (floats) 
        :param xy_int: 2d array of true xy positions (floats)
        :param config: dict with wfsim config
        :param resource: instance of the resource class

        returns two arrays of floats (mean drift time, drift time spread) 
        """
        drift_velocity_liquid = S2.get_avg_drift_velocity(z_int, xy_int, config, resource)
        if config['enable_field_dependencies']['diffusion_longitudinal_map']:
            diffusion_constant_longitudinal = resource.field_dependencies_map(z_int, xy_int, map_name='diffusion_longitudinal_map')  # cm²/s
            diffusion_constant_longitudinal *= 1e-9  # cm²/ns
        else:
            diffusion_constant_longitudinal = config['diffusion_constant_longitudinal']

        drift_time_mean = - z_int / \
            drift_velocity_liquid + config['drift_time_gate']
        drift_time_mean = np.clip(drift_time_mean, 0, np.inf)
        drift_time_spread = np.sqrt(2 * diffusion_constant_longitudinal * drift_time_mean)
        drift_time_spread /= drift_velocity_liquid
        return drift_time_mean, drift_time_spread
    
    @staticmethod
    def get_s2_light_yield(positions, config, resource):
        """Calculate s2 light yield...
        
        :param positions: 2d array of positions (floats)
        :param config: dict with wfsim config
        :param resource: instance of the resource class
        
        returns array of floats (mean expectation) 
        """
        
        if config.get('se_gain_from_map', False):
           sc_gain = resource.se_gain_map(positions)
        else:
           # calculate it from MC pattern map directly if no "se_gain_map" is given
           sc_gain = resource.s2_correction_map(positions)
           sc_gain *= config['s2_secondary_sc_gain']

       # depending on if you use the data driven or mc pattern map for light yield for S2
       # the shape of n_photon_hits will change. Mc needs a squeeze
        if len(sc_gain.shape) != 1:
            sc_gain=np.squeeze(sc_gain, axis=-1)

        # sc gain should has the unit of pe / electron, here we divide 1 + dpe to get nphoton / electron
        sc_gain /= 1 + config['p_double_pe_emision']

        # data driven map contains nan, will be set to 0 here
        sc_gain[np.isnan(sc_gain)] = 0

        return sc_gain

    @staticmethod
    def get_electron_yield(n_electron, xy_int, z_int, config, resource):
        """Drift electrons up to the gas interface and absorb them

        :param n_electron: 1d array with ints as number of electrons
        :param xy_int: 2d array of xy interaction positions (floats)
        :param z_int: 1d array of floats with the z interaction positions (floats)
        :param config: dict with wfsim config
        :param resource: instance of the resource class

        returns 1d array ints with number of electrons
        """
        # Average drift time of the electrons
        drift_time_mean, drift_time_spread = S2.get_s2_drift_time_params(z_int, xy_int, config, resource)

        # Absorb electrons during the drift
        electron_lifetime_correction = np.exp(- 1 * drift_time_mean /
                                              config['electron_lifetime_liquid'])

        if config.get('ext_eff_from_map', False):
            # Extraction efficiency is g2(x,y)/SE_gain(x,y)
            cy = config['g2_mean']*resource.s2_correction_map(xy_int)* \
                 electron_lifetime_correction/resource.se_gain_map(xy_int)
        else:
            cy = electron_lifetime_correction*config['electron_extraction_yield']

        # Remove electrons in insensitive volume
        if config['enable_field_dependencies']['survival_probability_map']:
            p_surv = resource.field_dependencies_map(z_int, xy_int, map_name='survival_probability_map')
            if np.any(p_surv<0) or np.any(p_surv>1):
                # FIXME: this is necessary due to map artefacts, such as negative or values >1
                p_surv=np.clip(p_surv, a_min = 0, a_max = 1)
            cy *= p_surv
        n_electron = np.random.binomial(n=n_electron, p=cy)
        
        return n_electron

    @staticmethod
    @njit
    def electron_timings(t, n_electron, drift_time_mean, drift_time_spread, sc_gain, timings, gains,
                         electron_trapping_time):
        """Calculate arrival times of the electrons. Data is written to the timings and gains arrays
        :param t: 1d array of ints
        :param n_electron:1 d array of ints
        :param drift_time_mean: 1d array of floats
        :param drift_time_spread: 1d array of floats
        :param sc_gain: secondary scintillation gain       
        :param timings: empty array with length sum(n_electron)
        :param gains: empty array with length sum(n_electron)
        :param electron_trapping_time: configuration values
        """
        assert len(timings) == np.sum(n_electron)
        assert len(gains) == np.sum(n_electron)
        assert len(sc_gain) == len(t)

        i_electron = 0
        for i in np.arange(len(t)):
            # Calculate electron arrival times in the ELR region
            for _ in np.arange(n_electron[i]):
                _timing = np.random.exponential(electron_trapping_time)
                _timing += np.random.normal(drift_time_mean[i], drift_time_spread[i])
                timings[i_electron] = t[i] + int(_timing)

                # add manual fluctuation to sc gain
                gains[i_electron] = sc_gain[i]
                i_electron += 1

    @staticmethod
    def get_n_photons(t, n_electron, z_int, xy_int, sc_gain, config, resource):
        """Generates photon timings for S2s.
        Returns a list of photon timings and instructions repeated for original electron
        
        :param t: 1d int array time of s2
        :param n_electron: 1d float array number of electrons to simulate
        :param z_int: float array. true Z interaction positions
        :param xy_int: 2d float array, true xy interaction positions
        :param sc_gain: float, secondary s2 gain
        :param config: dict of the wfsim config
        :param resource: instance of the resource class """
        # Get electron timings
        drift_time_mean, drift_time_spread = S2.get_s2_drift_time_params(z_int, xy_int, config, resource)
        _electron_timings = np.zeros(np.sum(n_electron), np.int64)
        _electron_gains = np.zeros(np.sum(n_electron), np.float64)
        S2.electron_timings(t, n_electron, drift_time_mean, drift_time_spread, sc_gain,
                            _electron_timings, _electron_gains, config['electron_trapping_time'])

        # Populate with photons per e/ per position
        n_photons_per_ele = np.random.poisson(_electron_gains)
        n_photons_per_ele += np.random.normal(0, config.get('s2_gain_spread', 0), len(n_photons_per_ele)).astype(np.int64)
        n_photons_per_ele[n_photons_per_ele < 0] = 0
        #
        n_photons_per_xy = np.cumsum(np.pad(n_photons_per_ele, [1, 0]))[np.cumsum(n_electron)]
        n_photons_per_xy = np.diff(np.pad(n_photons_per_xy, [1, 0]))

        return n_photons_per_xy, n_photons_per_ele, _electron_timings

    @staticmethod
    @njit
    def _luminescence_timings_simple(n, dG, E0, r, dr, rr, alpha, uE, p, n_photons):
        """
        Luminescence time distribution computation, calculates emission timings of photons from the excited electrons
        return 1d nested array with ints
        """
        emission_time = np.zeros(np.sum(n_photons), np.int64)

        ci = 0
        for i in range(n):
            npho = n_photons[i]
            dt = dr / (alpha * E0[i] * rr)
            dy = E0[i] * rr / uE - 0.8 * p  # arXiv:physics/0702142
            avgt = np.sum(np.cumsum(dt) * dy) / np.sum(dy)

            j = np.argmax(r <= dG[i])
            t = np.cumsum(dt[j:]) - avgt
            y = np.cumsum(dy[j:])

            probabilities = np.random.rand(npho)
            emission_time[ci:ci+npho] = np.interp(probabilities, y / y[-1], t).astype(np.int64)
            ci += npho

        return emission_time

    @staticmethod
    def luminescence_timings_simple(xy, n_photons, config, resource):
        """
        Luminescence time distribution computation according to simple s2 model (many many many single electrons)
        :param xy: 1d array with positions
        :param n_photons: 1d array with ints for number of xy positions
        :param config: dict wfsim config
        :param resource: instance of wfsim resource
        returns _luminescence_timings_simple
        """
        assert len(n_photons) == len(xy), 'Input number of n_photons should have same length as positions'

        number_density_gas = config['pressure'] / \
            (units.boltzmannConstant * config['temperature'])
        alpha = config['gas_drift_velocity_slope'] / number_density_gas
        uE = units.kV / units.cm
        pressure = config['pressure'] / units.bar

        if config.get('enable_gas_gap_warping', True):
            dG = resource.gas_gap_length(xy)
        else:
            dG = np.ones(len(xy)) * config['elr_gas_gap_length']
        rA = config['anode_field_domination_distance']
        rW = config['anode_wire_radius']
        dL = config['gate_to_anode_distance'] - dG

        VG = config['anode_voltage'] / (1 + dL / dG / config['lxe_dielectric_constant'])
        E0 = VG / ((dG - rA) / rA + np.log(rA / rW))  # V / cm

        dr = 0.0001  # cm
        r = np.arange(np.max(dG), rW, -dr)
        rr = np.clip(1 / r, 1 / rA, 1 / rW)

        return S2._luminescence_timings_simple(len(xy), dG, E0, 
                                               r, dr, rr, alpha, uE,
                                               pressure, n_photons)
    
    @staticmethod
    def luminescence_timings_garfield(xy, n_photons, config, resource, confine_position=None):
        """
        Luminescence time distribution computation according to garfield scintillation maps
        :param xy: 1d array with positions
        :param n_photons: 1d array with ints for number of xy positions
        :param config: dict wfsim config
        :param resource: instance of wfsim resource
        :param confine_position: if float, confine extraction region +/- this position around anode wires
        returns 2d array with ints for photon timings of input param 'shape'
        """
        assert 's2_luminescence' in resource.__dict__, 's2_luminescence model not found'
        assert len(n_photons) == len(xy), 'Input number of n_electron should have same length as positions'
        assert len(resource.s2_luminescence['t'].shape) == 2, 'Timing data is expected to have D2'

        if type(confine_position)==float:
            distance = np.random.uniform(-confine_position, confine_position, len(xy))
        else:
            tilt = config.get('anode_xaxis_angle', np.pi / 4)
            pitch = config.get('anode_pitch', 0.5)
            rotation_mat = np.array(((np.cos(tilt), -np.sin(tilt)), (np.sin(tilt), np.cos(tilt))))
            jagged = lambda relative_y: (relative_y + pitch / 2) % pitch - pitch / 2
            distance = jagged(np.matmul(xy, rotation_mat)[:, 1])  # shortest distance from any wire

        index_row = [np.argmin(np.abs(d - resource.s2_luminescence['x'])) for d in distance]
        index_row = np.repeat(index_row, n_photons).astype(np.int64)
        index_col = np.random.randint(0, resource.s2_luminescence['t'].shape[1], np.sum(n_photons), np.int64)
        
        avgt = np.average(resource.s2_luminescence['t']).astype(int)
        return resource.s2_luminescence['t'][index_row, index_col].astype(np.int64) - avgt
    
    @staticmethod
    @njit
    def draw_excitation_times(inv_cdf_list, hist_indices, nph, diff_nearest_gg, d_gas_gap):
        
        """
        Draws the excitation times from the GARFIELD electroluminescence map
        
        :param inv_cdf_list:       List of inverse CDFs for the excitation time histograms
        :param hist_indices:       The index of the histogram which refers to the gas gap
        :param nph:                A 1-d array of the number of photons per electron
        :param diff_nearest_gg:    The difference between the gas gap from the map
                                   (continuous value) and the nearest (discrete) value
                                   of the gas gap corresponding to the excitation time
                                   histograms
        :param d_gas_gap:          Spacing between two consecutive gas gap values
        
        returns time of each photon
        """
        
        inv_cdf_len = len(inv_cdf_list[0])
        timings = np.zeros(np.sum(nph))
        upper_hist_ind = np.clip(hist_indices+1, 0, len(inv_cdf_list)-1)
        
        count = 0
        for i, (hist_ind, u_hist_ind, n, dngg) in enumerate(zip(hist_indices,
                                                                upper_hist_ind, 
                                                                nph,
                                                                diff_nearest_gg)):
            
            #There are only 10 values of gas gap separated by 0.1mm, so we interpolate
            #between two histograms
            
            interp_cdf = ((inv_cdf_list[u_hist_ind]-inv_cdf_list[hist_ind])*(dngg/d_gas_gap)
                           +inv_cdf_list[hist_ind])

            samples = np.random.uniform(0, inv_cdf_len-1, n)
            t1 = interp_cdf[np.floor(samples).astype('int')]
            t2 = interp_cdf[np.ceil(samples).astype('int')]
            T = (t2-t1)*(samples - np.floor(samples))+t1
            if n!=0:
                T = T-np.mean(T)
            
            #subtract mean to get proper drift time and z correlation
            timings[count:count+n] = T
            count+=n
        return timings
    
    @staticmethod
    def luminescence_timings_garfield_gasgap(xy, n_photons, resource):
        """
        Luminescence time distribution computation according to garfield scintillation maps
        which are ONLY drawn from below the anode, and at different gas gaps
        :param xy: 1d array with positions
        :param n_photons: 1d array with ints for number of xy positions
        :param resource: instance of wfsim resource

        returns 2d array with ints for photon timings of input param 'shape'
        """
        assert 's2_luminescence_gg' in resource.__dict__, 's2_luminescence_gg model not found'
        assert len(n_photons) == len(xy), 'Input number of n_electron should have same length as positions'
        
        d_gasgap = resource.s2_luminescence_gg['gas_gap'][1]-resource.s2_luminescence_gg['gas_gap'][0]
        
        cont_gas_gaps = resource.garfield_gas_gap_map(xy)
        draw_index = np.digitize(cont_gas_gaps, resource.s2_luminescence_gg['gas_gap'])-1
        diff_nearest_gg = cont_gas_gaps - resource.s2_luminescence_gg['gas_gap'][draw_index]
        
        return S2.draw_excitation_times(resource.s2_luminescence_gg['timing_inv_cdf'],
                                        draw_index,
                                        n_photons,
                                        diff_nearest_gg,
                                        d_gasgap)
    
    @staticmethod
    def optical_propagation(channels, config, spline):
        """Function getting times from s2 timing splines:
        :param channels: The channels of all s2 photon
        :param config: current configuration of wfsim
        :param spline: pointer to s2 optical propagation splines from resources
        """
        prop_time = np.zeros_like(channels)
        u_rand = np.random.rand(len(channels))[:, None]

        is_top = channels < config['n_top_pmts']
        prop_time[is_top] = spline(u_rand[is_top], map_name='top')

        is_bottom = channels >= config['n_top_pmts']
        prop_time[is_bottom] = spline(u_rand[is_bottom], map_name='bottom')

        return prop_time.astype(np.int64)

    @staticmethod
    def photon_timings(positions, n_photons_per_xy, _electron_timings, n_photons_per_ele, _photon_channels, phase, config, resource):
        """Get photon times and add delays based on models

        :param positions: 2d float array, xy positions of s2
        :param n_photons_per_xy: number of photons for each xy position
        :param _electron_timings: electron timings
        :param n_photons_per_ele: number of photons for each electron
        :param _photon_channels: channel of each photon
        :param config: dict of the wfsim config
        :param phase: gas
        :param resource: instance of the resource class
        """

        # Luminescence Timings
        if config['s2_luminescence_model']=='simple':
            _photon_timings = S2.luminescence_timings_simple(positions, n_photons_per_xy,
                                                             config=config,
                                                             resource=resource)
        elif config['s2_luminescence_model']=='garfield':
            confine_position=None
            if 's2_garfield_confine_position' in config:
                if config['s2_garfield_confine_position'] > 0.0:
                    confine_position=config['s2_garfield_confine_position']
            _photon_timings = S2.luminescence_timings_garfield(positions, n_photons_per_xy,
                                                               config=config,
                                                               resource=resource,
                                                               confine_position=confine_position)
            
        elif config['s2_luminescence_model']=='garfield_gas_gap':
            _photon_timings = S2.luminescence_timings_garfield_gasgap(positions, n_photons_per_xy,
                                                                      resource=resource)
        else:
            raise KeyError(f"{config['s2_luminescence_model']} is not valid! Use 'simple' or 'garfield' or 'garfield_gas_gap'")

        # Emission Delay
        _photon_timings += Pulse.singlet_triplet_delays(len(_photon_timings), config['singlet_fraction_gas'], config, phase)

        # Optical Propagation Delay
        if "optical_propagation" in config['s2_time_model']:
            # optical propagation splitting top and bottom
            _photon_timings += S2.optical_propagation(_photon_channels, config, resource.s2_optical_propagation_spline)
        elif "zero_delay" in config['s2_time_model']:
            # no optical propagation delay
            _photon_timings += np.zeros_like(_photon_timings, dtype=np.int64)
        elif "s2_time_spread around zero" in config['s2_time_model']:
            # simple/existing delay
            _photon_timings += np.random.normal(0, config['s2_time_spread'], len(_photon_timings)).astype(np.int64)
        else:
            raise KeyError(f"{config['s2_time_model']} is not in any of the valid s2 time models")

        # repeat for n photons per electron # Should this be before adding delays?
        _photon_timings += np.repeat(_electron_timings, n_photons_per_ele)

        return _photon_timings

    @staticmethod
    def s2_pattern_map_diffuse(n_electron, z, xy, config, resource):
        """Returns an array of pattern of shape [n interaction, n PMTs]
        pattern of each interaction is an average of n_electron patterns evaluated at
        diffused position near xy. The diffused positions sample from 2d symmetric gaussian
        with spread scale with sqrt of drift time.

        :param n_electron: a 1d int array
        :param z: a 1d float array
        :param xy: a 2d float array of shape [n interaction, 2]
        :param config: dict of the wfsim config
        :param resource: instance of the resource class
        """
        assert all(z < 0), 'All S2 in liquid should have z < 0'
        drift_velocity_liquid=S2.get_avg_drift_velocity(z, xy, config, resource)

        if config['enable_field_dependencies']['diffusion_transverse_map']:
            diffusion_constant_radial = resource.field_dependencies_map(z, xy, map_name='diffusion_radial_map')  # cm²/s
            diffusion_constant_azimuthal = resource.field_dependencies_map(z, xy, map_name='diffusion_azimuthal_map') # cm²/s
            diffusion_constant_radial *= 1e-9  # cm²/ns
            diffusion_constant_azimuthal *= 1e-9  # cm²/ns
        else:
            diffusion_constant_transverse = getattr(config, 'diffusion_constant_transverse', 0)
            diffusion_constant_radial = diffusion_constant_transverse
            diffusion_constant_azimuthal = diffusion_constant_transverse

        drift_time_mean = - z / drift_velocity_liquid
        hdiff_stdev_radial = np.sqrt(2 * diffusion_constant_radial * drift_time_mean)
        hdiff_stdev_azimuthal = np.sqrt(2 * diffusion_constant_azimuthal * drift_time_mean)
        
        hdiff_radial = np.random.normal(0, 1, np.sum(n_electron)) * np.repeat(hdiff_stdev_radial, n_electron, axis=0)
        hdiff_azimuthal = np.random.normal(0, 1, np.sum(n_electron)) * np.repeat(hdiff_stdev_azimuthal, n_electron, axis=0)
        hdiff = np.column_stack([hdiff_radial, hdiff_azimuthal])
        theta = np.arctan2(xy[:,1], xy[:,0])
        matrix = np.array([[np.cos(theta), np.sin(theta)], [-np.sin(theta), np.cos(theta)]]).T
        hdiff = np.vstack([(matrix[i] @ np.split(hdiff, np.cumsum(n_electron))[:-1][i].T).T for i in range(len(matrix))])
        # Should we also output this xy position in truth?
        xy_multi = np.repeat(xy, n_electron, axis=0) + hdiff  # One entry xy per electron
        # Remove points outside tpc, and the pattern will be the average inside tpc
        # Should be done naturally with the s2 pattern map, however, there's some bug there, so we apply this hard cut
        mask = np.sum(xy_multi ** 2, axis=1) <= config['tpc_radius'] ** 2

        if isinstance(resource.s2_pattern_map, DummyMap):
            output_dim = resource.s2_pattern_map.shape[-1]
        else:
            output_dim = resource.s2_pattern_map.data['map'].shape[-1]
        pattern = np.zeros((len(n_electron), output_dim))
        n0 = 0
        # Average over electrons for each s2
        for ix, ne in enumerate(n_electron):
            s = slice(n0, n0+ne)
            pattern[ix, :] = np.average(resource.s2_pattern_map(xy_multi[s][mask[s]]), axis=0)
            n0 += ne

        return pattern

    @staticmethod
    def photon_channels(n_electron, z_obs, positions, _instruction, config, resource):
        """Set the _photon_channels property list of length same as _photon_timings
        :param n_electron: a 1d int array
        :param z_obs: a 1d float array
        :param positions: a 2d float array of shape [n interaction, 2] for the xy coordinate
        :param _instruction: array of instructions with dtype wfsim.instructions_dtype
        :param config: dict wfsim config
        :param resource: instance of resource class
        """
        if len(_instruction) == 0:
            _photon_channels = np.zeros(0, dtype=np.int64)
            return _photon_channels

        #aft = config['s2_mean_area_fraction_top']
        aft_sigma = config.get('s2_aft_sigma', 0.0)
        aft_skewness = config.get('s2_aft_skewness', 0.0)

        channels = np.arange(config['n_tpc_pmts']).astype(np.int64)
        top_index = np.arange(config['n_top_pmts'])
        bottom_index = np.array(config['channels_bottom'])

        if config.get('diffusion_constant_transverse', 0) > 0:
            pattern = S2.s2_pattern_map_diffuse(n_electron, z_obs, positions, config, resource)  # [position, pmt]
        else:
            pattern = resource.s2_pattern_map(positions)  # [position, pmt]

        if pattern.shape[1] - 1 not in bottom_index:
            pattern = np.pad(pattern, [[0, 0], [0, len(bottom_index)]], 
                             'constant', constant_values=1)

        # Remove turned off pmts
        pattern[:, np.in1d(channels, config['turned_off_pmts'])] = 0

        sum_pat = np.sum(pattern, axis=1).reshape(-1, 1)
        pattern = np.divide(pattern, sum_pat, out=np.zeros_like(pattern), where=sum_pat != 0)

        assert pattern.shape[0] == len(positions)
        assert pattern.shape[1] == len(channels)

        _buffer_photon_channels = []
        # Randomly assign to channel given probability of each channel
        for unique_i, count in zip(*np.unique(_instruction, return_counts=True)):
            pat = pattern[unique_i]  # [pmt]

            if aft_sigma != 0:  # Redistribute pattern with user specified aft smearing
                _cur_aft=np.sum(pat[top_index])/np.sum(pat)
                _new_aft=_cur_aft*skewnorm.rvs(loc=1.0, scale=aft_sigma, a=aft_skewness)
                _new_aft=np.clip(_new_aft, 0, 1)
                pat[top_index]*=(_new_aft/_cur_aft)
                pat[bottom_index]*=(1 - _new_aft)/(1 - _cur_aft)
                # old code from Peter
                #pat[top_index] = pat[top_index] / pat[top_index].sum() * _aft
                #pat[bottom_index] = pat[bottom_index] / pat[bottom_index].sum() * (1 - _aft)

            if np.isnan(pat).sum() > 0:  # Pattern map return zeros
                _photon_channels = np.array([-1] * count)
            else:
                _photon_channels = np.random.choice(
                    channels,
                    size=count,
                    p=pat,
                    replace=True)

            _buffer_photon_channels.append(_photon_channels)

        _photon_channels = np.concatenate(_buffer_photon_channels)
        return _photon_channels
