# -*- coding: utf-8 -*-
# Copyright (C) Scott Coughlin (2017)
#
# This file is part of cosmic.
#
# cosmic is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# cosmic is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with cosmic.  If not, see <http://www.gnu.org/licenses/>.

"""`gxreal`
"""

import numpy as np
from gwpy.utils import mp as mp_utils
import pandas as pd
import cosmic.MC_samp as MC_sample
import cosmic.utils as utils
import scipy.stats as stats
import cosmic.GW_calcs as GW_calcs

__author__ = 'Katelyn Breivik <katie.breivik@gmail.com>'
__credits__ = 'Scott Coughlin <scott.coughlin@ligo.org>'
__all__ = ['GxReal']

G = 6.67384e-11
c = 2.99792458e8
parsec = 3.08567758e16
Rsun = 6.955e8
Msun = 1.9891e30
day = 86400.0
rsun_in_au = 215.0954
day_in_year = 365.242
sec_in_day = 86400.0
sec_in_hour = 3600.0
hrs_in_day = 24.0
sec_in_year = 3.15569e7

class GxReal(object):
    def __init__(self, fixed_pop, m_tot_samp, gx_model, gx_component, dat_list):
        '''
        Initialize GxReal
        '''
        self.fixed_pop = fixed_pop
        self.fixed_mass = m_tot_samp
        self.gx_model = gx_model
        self.gx_component = gx_component
        self.dat_list = dat_list 

        return


    def compute_n_sample(self):
        """Use the fixed population and total mass sampled to generate the 
        fixed population to compute the number of systems in the Milky Way
        realization

        Parameters
        ----------
        fixedpop : DataFrame
            Contains binary parameters from an evolved population
            generated by the runFixedPop executable
        fixed_mass : float
            Total mass sampled to generate the fixed population
        gx_component : str
            Milky Way component for which we are generating the population
            realization; choose from 'ThinDisk', 'ThickDisk', 'Bulge'

        Returns
        -------
        n_samp : int
            The number of systems in the Milky Way realization
        """

        # Compute the mass weighted number of systems 
        # in a realistic Milky Way pop
        #######################################################################
        component_mass = MC_sample.select_component_mass(self.gx_component)
        n_samp = MC_sample.mass_weighted_number(self.fixed_pop, self.fixed_mass, component_mass)
                 
        return n_samp


    def sample_population(self):
        """Once the fixed population is evolved, we Monte-Carlo
        sample the binary parameters to generate a Milky Way realization.

        Parameters
        ----------
        fixedpop : DataFrame 
            Contains binary parameters from an evolved population
            generated by the runFixedPop executable
        n_samp : int
            The number of systems in the Milky Way realization
        gx_component : str
            Milky Way component for which we are generating the population 
            realization; choose from 'ThinDisk', 'ThickDisk', 'Bulge'
        gx_model : str
            Model for spatial distribution of binaries in the Galactic 
            component; Default='McMillan'       
        dat_list : list
            List containing the parameters to MC sample to generate the 
            Galactic realization

        Returns
        -------
        realization : DataFrame
            Milky Way population realization of size n_samp
        """

        if not hasattr(self, 'n_samp'):
            self.n_samp = self.compute_n_sample()

        # Based on the user supplied filter flags, filter the
        # population to reduce the sample to only the relevant population
        #######################################################################
        # Transform the fixed population to have limits between 0 and 1
        # then transform to logit space to maintain the population boundaries
        # and create a KDE using knuth's rule to select the bandwidth
        #######################################################################
        dat_kde = utils.dat_transform(self.fixed_pop, self.dat_list)   
        bw = utils.knuth_bw_selector(dat_kde)
        dat_kernel = stats.gaussian_kde(dat_kde, bw_method=bw)
   
        # Sample from the KDE
        #######################################################################
        binary_dat_trans = dat_kernel.resample(self.n_samp)

        binary_dat = utils.dat_un_transform(binary_dat_trans, self.fixed_pop, self.dat_list)
        
        # Sample positions and orientations for each sampled binary
        #######################################################################
        xGx, yGx, zGx, inc, OMEGA, omega = MC_sample.galactic_positions(self.gx_component,
                                                        size = len(binary_dat[0]),
                                                        model=self.gx_model)
        
        binary_sample_positions = np.vstack([xGx, yGx, zGx, inc, OMEGA, omega]) 
        # Create a single DataFrame for the Galactic realization
        #######################################################################
        if len(self.dat_list) < 4:
            full_sample = np.vstack([binary_dat,np.zeros(self.n_samp),binary_sample_positions]).T
            column_list = self.dat_list + ['ecc', 'xGx', 'yGx', 'zGx', 'inc', 'OMEGA', 'omega']
        else:
            full_sample = np.concatenate([binary_dat,binary_sample_positions]).T
            column_list = self.dat_list + ['xGx', 'yGx', 'zGx', 'inc', 'OMEGA', 'omega']
        realization = pd.DataFrame(full_sample,\
                                   columns = column_list) 

        return realization

    def LISA_obs(self, T_obs):
        """Computes the gravitational wave signal from the population
        that will be observable by LISA, including SNR and PSD according
        to the user input

        Parameters
        ----------
        realization : DataFrame
            Milky Way population realization of size n_samp
        T_obs : float
            LISA observation time in seconds

        Returns
        -------
        PSD_dat : DataFrame
            DataFrame containing the power spectral density for all systems
            in realization 
        """
        # Compute the PSD
        #######################################################################
        PSD_dat = GW_calcs.LISA_PSD(self.realization.mass_1,
                                     self.realization.mass_2,
                                     10**self.realization.porb,
                                     self.realization.ecc,
                                     self.realization.xGx,
                                     self.realization.yGx,
                                     self.realization.zGx,
                                     150, T_obs)
        return PSD_dat

    def LISA_full_obs(self, T_obs):
        """Computes the gravitational wave signal from the population
        that will be observable by LISA, including SNR and PSD according
        to the user input

        Parameters
        ----------
        realization : DataFrame
            Milky Way population realization of size n_samp
        T_obs : float
            LISA observation time in seconds

        Returns
        -------
        SNR_dat : DataFrame
            DataFrame containing signal to noise ratios for all systems 
            in realization
        PSD_dat : DataFrame
            DataFrame containing the power spectral density for all systems 
            in realization 
        foreground_dat : DataFrame
            DataFrame containing power spectral density of the population
            over a set of frequency bins with width set by the LISA 
            observation time           
        """

        # Compute the PSD
        #######################################################################
        PSD_dat = GW_calcs.LISA_PSD(self.realization.mass_1, 
                                     self.realization.mass_2,
                                     10**self.realization.porb,
                                     self.realization.ecc, 
                                     self.realization.xGx,
                                     self.realization.yGx,     
                                     self.realization.zGx,
                                     150, T_obs)
        # Compute the foreground from the PSD dataframe
        #######################################################################
        foreground_dat = GW_calcs.compute_foreground(PSD_dat, T_obs)
        
        LISA_plus_foreground = GW_calcs.LISA_combo(foreground_dat) 
         
        # Compute the SNR
        ####################################################################### 
        SNR_dat = GW_calcs.LISA_SNR(self.realization.mass_1,
                                     self.realization.mass_2,
                                     10**self.realization.porb,
                                     self.realization.ecc,
                                     self.realization.xGx,
                                     self.realization.yGx,
                                     self.realization.zGx,
                                     150, T_obs, 
                                     foreground_dat,
                                     LISA_plus_foreground) 

        return SNR_dat, PSD_dat, foreground_dat
         
