import logging
import numpy as np

from pytplot import get_data

from pyspedas.utilities.interpol import interpol
from pyspedas.particles.spd_part_products.spd_pgs_make_tplot import spd_pgs_make_tplot
from pyspedas.particles.spd_part_products.spd_pgs_limit_range import spd_pgs_limit_range
from pyspedas.particles.spd_part_products.spd_pgs_progress_update import spd_pgs_progress_update
from pyspedas.particles.spd_part_products.spd_pgs_do_fac import spd_pgs_do_fac
from pyspedas.particles.spd_part_products.spd_pgs_regrid import spd_pgs_regrid
from pyspedas.particles.moments.spd_pgs_moments import spd_pgs_moments
from pyspedas.particles.moments.spd_pgs_moments_tplot import spd_pgs_moments_tplot

from pyspedas.mms.fpi.mms_get_fpi_dist import mms_get_fpi_dist
from pyspedas.mms.hpca.mms_get_hpca_dist import mms_get_hpca_dist
from pyspedas.mms.particles.mms_convert_flux_units import mms_convert_flux_units
from pyspedas.mms.particles.mms_pgs_clean_data import mms_pgs_clean_data
from pyspedas.mms.particles.mms_pgs_clean_support import mms_pgs_clean_support
from pyspedas.mms.particles.mms_pgs_make_fac import mms_pgs_make_fac
from pyspedas.mms.particles.mms_pgs_split_hpca import mms_pgs_split_hpca
from pyspedas.mms.particles.mms_pgs_make_e_spec import mms_pgs_make_e_spec
from pyspedas.mms.particles.mms_pgs_make_phi_spec import mms_pgs_make_phi_spec
from pyspedas.mms.particles.mms_pgs_make_theta_spec import mms_pgs_make_theta_spec
from pyspedas.mms.particles.mms_part_des_photoelectrons import mms_part_des_photoelectrons

logging.captureWarnings(True)
logging.basicConfig(format='%(asctime)s: %(message)s', datefmt='%d-%b-%y %H:%M:%S', level=logging.INFO)

def mms_part_products(in_tvarname, units='eflux', species='e', data_rate='fast', instrument='fpi', probe='1',
    output=['energy', 'theta', 'phi'], energy=None, phi=None, theta=None, pitch=None, gyro=None, mag_name=None,
    pos_name=None, fac_type='mphigeo', sc_pot_name=None, correct_photoelectrons=False, zero_negative_values=False,
    internal_photoelectron_corrections=False, disable_photoelectron_corrections=False, no_regrid=False,
    regrid=[32, 16]):
    """
    Generate spectra and moments from 3D MMS particle data; note: this routine isn't
    meant to be called directly - see the wrapper mms_part_getspec instead.

    Input
    ----------
        in_tvarname: str
            Name of the tplot variable containing MMS 3D particle distribution data

    Parameters
    ----------
        units: str
            Specify units of output variables; must be 'eflux' to calculate moments

            valid options:
            'flux'   -  # / (cm^2 * s * sr * eV)
            'eflux'  -  eV / (cm^2 * s * sr * eV)  <default>
            'df_cm'  -  s^3 / cm^6
            'df_km'  -  s^3 / km^6

        species: str
            Specify the species of the input tplot variable

        data_rate: str
            Data rate of the input data

        instrument: str
            Instrument (FPI or HPCA)

        probe: int or str
            Spacecraft probe #

        output: str or list of str
            Output variables; options: 
                'energy': energy spectrograms
                'theta': theta spectrograms
                'phi': phi spectrograms
                'pa': pitch-angle spectrograms
                'gyro': gyro-phase spectrograms
                'moments': plasma moments

        energy: list of float
            Energy range [min, max], in eV

        phi: list of float
            Phi range [min, max], in degrees

        theta: list of float
            Theta range [min, max], in degrees

        pitch: list of float
            Pitch-angle range [min, max], in degrees

        gyro: list of float
            Gyro-angle range [min, max], in degrees

        mag_name: str
            Tplot variable containing magnetic field data for
            moments and FAC transformations

        pos_name: str
            Tplot variable containing spacecraft position for
            FAC transformations

        sc_pot_name: str
            Tplot variable containing spacecraft potential data
            for moments corrections

        fac_type: str
            Field aligned coordinate system variant; default: 'mphigeo'
            options: 'phigeo', 'mphigeo', 'xgse'

        correct_photoelectrons: bool
            Flag to correct FPI data for photoelectrons 
            (defaults to True for FPI electron data - disable with the parameter below)

        disable_photoelectron_corrections: bool
            Flag to disable FPI photoelectron corrections

        internal_photoelectron_corrections: bool
            Apply internal photoelectron corrections

        zero_negative_values: bool
            Turn negative values to 0 after doing the photoelectron corrections (DES)
            
    Returns
    ----------
        Creates tplot variables containing spectrograms and moments

    """

    units = units.lower()
    if isinstance(probe, int):
        probe = str(probe)

    data_in = get_data(in_tvarname)

    if data_in is None:
        logging.error('Error, could not find the variable: ' + in_tvarname)
        return None

    if isinstance(output, str):
        output = output.split(' ')

    if instrument == 'fpi':
        dist_in = mms_get_fpi_dist(in_tvarname, species=species, probe=probe, data_rate=data_rate)
    elif instrument == 'hpca':
        dist_in = mms_get_hpca_dist(in_tvarname, species=species, probe=probe, data_rate=data_rate)
    else:
        logging.error('Error, unknown instrument: ' + instrument + '; valid options: fpi, hpca')
        return

    if instrument == 'hpca':
        data_times = mms_get_hpca_dist(in_tvarname, species=species, probe=probe, data_rate=data_rate, times=True)
    else:
        data_times = data_in.times
    ntimes = len(data_times)

    # create rotation matrix to field aligned coordinates if needed
    fac_outputs = ['pa', 'gyro', 'fac_energy', 'fac_moments']
    fac_requested = len(set(output).intersection(fac_outputs)) > 0
    if fac_requested:
        fac_matrix = mms_pgs_make_fac(data_times, mag_name, pos_name, fac_type=fac_type)

        if fac_matrix is None:
            # problem creating the FAC matrices
            fac_requested = False

    out_energy = np.zeros((dist_in['n_times'], dist_in['n_energy']))
    out_energy_y = np.zeros((dist_in['n_times'], dist_in['n_energy']))
    out_theta = np.zeros((dist_in['n_times'], dist_in['n_theta']))
    out_phi = np.zeros((dist_in['n_times'], dist_in['n_phi']))
    out_theta_y = np.zeros((dist_in['n_times'], dist_in['n_theta']))
    out_phi_y = np.zeros((dist_in['n_times'], dist_in['n_phi']))
    if fac_requested:
        out_pad = np.zeros((dist_in['n_times'], dist_in['n_theta']))
        out_pad_y = np.zeros((dist_in['n_times'], dist_in['n_theta']))
        out_gyro = np.zeros((dist_in['n_times'], dist_in['n_phi']))
        out_gyro_y = np.zeros((dist_in['n_times'], dist_in['n_phi']))

    # moments
    if 'moments' in output:
        out_density = np.zeros(dist_in['n_times'])
        out_avgtemp = np.zeros(dist_in['n_times'])
        out_vthermal = np.zeros(dist_in['n_times'])
        out_flux = np.zeros([dist_in['n_times'], 3])
        out_velocity = np.zeros([dist_in['n_times'], 3])
        out_mftens = np.zeros([dist_in['n_times'], 6])
        out_ptens = np.zeros([dist_in['n_times'], 6])
        #out_ttens = np.zeros([dist_in['n_times'], 3, 3])

    out_vars = []
    last_update_time = None

    if 'moments' in output or correct_photoelectrons or internal_photoelectron_corrections:
        support_data = mms_pgs_clean_support(data_times, mag_name=mag_name, vel_name=None, sc_pot_name=sc_pot_name)
        mag_data = support_data[0]
        scpot_data = support_data[2]

    if disable_photoelectron_corrections:
        correct_photoelectrons = False
        internal_photoelectron_corrections = False

    # grab the DES photoelectron model if needed
    if (instrument != 'fpi' or species != 'e') and (correct_photoelectrons or internal_photoelectron_corrections):
        logging.error('Photoelectron corrections only valid for DES; no corrections will be applied.')
        correct_photoelectrons = False
        internal_photoelectron_corrections = False

    if correct_photoelectrons or internal_photoelectron_corrections:
        fpi_photoelectrons = mms_part_des_photoelectrons(in_tvarname)

        if fpi_photoelectrons is None:
            logging.error('Photoelectron model missing for this date; re-run without photoelectron corrections')
            return

        # will need stepper parities for burst mode data
        if data_rate == 'brst':
            parity = get_data('mms'+probe+'_des_steptable_parity_brst')

        startdelphi = get_data('mms'+probe+'_des_startdelphi_count_'+data_rate)

    for i in range(0, ntimes):
        last_update_time = spd_pgs_progress_update(last_update_time=last_update_time, current_sample=i, total_samples=ntimes, type_string=in_tvarname)

        if instrument == 'fpi':
            dist_in = mms_get_fpi_dist(in_tvarname, index=i, species=species, probe=probe, data_rate=data_rate)
        elif instrument == 'hpca':
            dist_in = mms_get_hpca_dist(in_tvarname, index=i, species=species, probe=probe, data_rate=data_rate)

        # apply the DES photoelectron corrections
        if correct_photoelectrons or internal_photoelectron_corrections:
            # From Dan Gershman's release notes on the FPI photoelectron model:
            # Find the index I in the startdelphi_counts_brst or startdelphi_counts_fast array
            # [360 possibilities] whose corresponding value is closest to th = e measured
            # startdelphi_count_brst or startdelphi_count_fast for the skymap of interest. The
            # closest index can be approximated by I = floor(startdelphi_count_brst/16) or I =
            # floor(startdelphi_count_fast/16)
            startdelphi_I = int(np.floor(startdelphi.y[i]/16.0))

            if data_rate == 'brst':
                parity_num = str(int(np.fix(parity.y[i])))

                bg_dist = fpi_photoelectrons['bgdist_p'+parity_num]
                n_value = fpi_photoelectrons['n_'+parity_num]

                fphoto = bg_dist.y[startdelphi_I, :, :, :]

                # need to interpolate using SC potential data to get Nphoto value
                nphoto_scpot_dependent = n_value.y[startdelphi_I, :]
                nphoto = interpol(nphoto_scpot_dependent, n_value.v, scpot_data[i])
            else:
                fphoto = fpi_photoelectrons['bg_dist'].y[startdelphi_I, :, :, :]

                # need to interpolate using SC potential data to get Nphoto value
                nphoto_scpot_dependent = fpi_photoelectrons['n'].y[startdelphi_I, :]
                nphoto = interpol(nphoto_scpot_dependent, fpi_photoelectrons['n'].v, scpot_data[i])

            # now, the corrected distribution function is simply f_corrected = f-fphoto*nphoto
            # note: transpose is to shuffle fphoto*nphoto to energy-azimuth-elevation, to match dist.data
            correction = fphoto*nphoto
            corrected_df = dist_in['data']-correction.transpose([2, 0, 1])

            if zero_negative_values:
                corrected_df[corrected_df < 0] = 0.0

            dist_in['data'] = corrected_df

        data = mms_convert_flux_units(dist_in, units=units)

        clean_data = mms_pgs_clean_data(data)

        # split hpca angle bins to be equal width in phi/theta
        # this is needed when skipping the regrid step
        if instrument == 'hpca':
            clean_data = mms_pgs_split_hpca(clean_data)

        # Apply phi, theta, & energy limits
        if energy is not None or theta is not None or phi is not None:
            clean_data = spd_pgs_limit_range(clean_data, energy=energy, theta=theta, phi=phi)

        # Build energy spectrogram
        if 'energy' in output:
            out_energy_y[i, :], out_energy[i, :] = mms_pgs_make_e_spec(clean_data)

        # Build theta spectrogram
        if 'theta' in output:
            out_theta_y[i, :], out_theta[i, :] = mms_pgs_make_theta_spec(clean_data, resolution=dist_in['n_theta'])

        # Build phi spectrogram
        if 'phi' in output:
            out_phi_y[i, :], out_phi[i, :] = mms_pgs_make_phi_spec(clean_data, resolution=dist_in['n_phi'])

        # Calculate the moments
        if 'moments' in output:
            if scpot_data is not None:
                scpot_val = scpot_data[i]
            else:
                scpot_val = 0.0

            moments = spd_pgs_moments(clean_data, sc_pot=scpot_val)
            out_density[i] = moments['density']
            out_avgtemp[i] = moments['avgtemp']
            out_vthermal[i] = moments['vthermal']
            out_flux[i, :] = moments['flux']
            out_velocity[i, :] = moments['velocity']
            out_mftens[i, :] = moments['mftens']
            out_ptens[i, :] = moments['ptens']

        # Perform transformation to FAC, regrid data, and apply limits in new coords
        if fac_requested:
            fac_data = spd_pgs_do_fac(clean_data, fac_matrix[i, :, :])

            if no_regrid == False:
                fac_data = spd_pgs_regrid(fac_data, regrid)

            fac_data['theta'] = 90.0-fac_data['theta']
            fac_data = spd_pgs_limit_range(fac_data, theta=pitch, phi=gyro)

        if 'pa' in output:
            out_pad_y[i, :], out_pad[i, :] = mms_pgs_make_theta_spec(fac_data, colatitude=True, resolution=dist_in['n_theta'])

        if 'gyro' in output:
            out_gyro_y[i, :], out_gyro[i, :] = mms_pgs_make_phi_spec(fac_data, resolution=dist_in['n_phi'])


    if 'moments' in output:
        # put all of the moments arrays into a hash table prior to passing to the tplot routine
        moments = {'density': out_density, 
              'flux': out_flux, 
              'mftens': out_mftens, 
              'velocity': out_velocity, 
              'ptens': out_ptens,
              'vthermal': out_vthermal,
              'avgtemp': out_avgtemp}
        moments_vars = spd_pgs_moments_tplot(moments, x=data_times, prefix=in_tvarname)
        out_vars.extend(moments_vars)

    if 'energy' in output:
        spd_pgs_make_tplot(in_tvarname+'_energy', x=data_times, y=out_energy_y, z=out_energy, units=units, ylog=True, ytitle=dist_in['data_name'] + ' \\ energy (eV)')
        out_vars.append(in_tvarname+'_energy')

    if 'theta' in output:
        spd_pgs_make_tplot(in_tvarname+'_theta', x=data_times, y=out_theta_y, z=out_theta, units=units, ytitle=dist_in['data_name'] + ' \\ theta (deg)')
        out_vars.append(in_tvarname+'_theta')

    if 'phi' in output:
        spd_pgs_make_tplot(in_tvarname+'_phi', x=data_times, y=out_phi_y, z=out_phi, units=units, ytitle=dist_in['data_name'] + ' \\ phi (deg)')
        out_vars.append(in_tvarname+'_phi')

    if 'pa' in output:
        spd_pgs_make_tplot(in_tvarname+'_pa', x=data_times, y=out_pad_y, z=out_pad, units=units, ytitle=dist_in['data_name'] + ' \\ PA (deg)')
        out_vars.append(in_tvarname+'_pa')

    if 'gyro' in output:
        spd_pgs_make_tplot(in_tvarname+'_gyro', x=data_times, y=out_gyro_y, z=out_gyro, units=units, ytitle=dist_in['data_name'] + ' \\ gyro (deg)')
        out_vars.append(in_tvarname+'_gyro')

    return out_vars