# covariance.py
#
# Create the first row of the covariance matrix for various noise models:
#  1) power-law noise using Eq. (7) of Bos et al. (2008).
#  2) white noise
#
# Bos, MS, Fernandes, RMS, Williams, SDP & Bastos, L (2008). "Fast error 
# analysis of continuous GPS observations". Journal of Geodesy, 82(3), 157-166.
#
#
# This file is part of HectorP 0.1.1.
#
# HectorP 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.
#
# HectorP 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 
# HectorP. If not, see <https://www.gnu.org/licenses/>.
#
# 18/ 5/2020 Machiel Bos, Santa Clara
# 29/12/2021 Machiel Bos, Santa Clara
#==============================================================================

import numpy as np
import sys
import math
from hectorp.white import White
from hectorp.ggm import GGM
from hectorp.powerlaw import Powerlaw
from hectorp.varyingannual import VaryingAnnual
from hectorp.ar1 import AR1
from hectorp.matern import Matern
from hectorp.control import Control
from hectorp.control import SingletonMeta

#==============================================================================
# Subroutines
#==============================================================================


class Covariance(metaclass=SingletonMeta):


    def __init__(self):
        """ initialise class
        """

        #--- Get control parameters
        control = Control()
        try:
            self.verbose = control.params['Verbose']
        except:
            self.verbose = True

        try:
            self.noisemodel_names = control.params['NoiseModels']
            if isinstance(self.noisemodel_names,list)==False:
                self.noisemodel_names = [self.noisemodel_names]
        except Exception as e:
            print(e)
            sys.exit()

        #--- Create list of noise model class instances
        self.noisemodels = []
        i = 0
        for noisemodel_name in self.noisemodel_names:
            if noisemodel_name in ['FlickerGGM','RandomWalkGGM','GGM']:
                ClassName = getattr(sys.modules[__name__], 'GGM')
                if noisemodel_name=='FlickerGGM':
                    class_ = ClassName(0.5)
                elif noisemodel_name=='RandomWalkGGM':
                    class_ = ClassName(1.0)
                else: 
                    try:
                        d = -0.5*control.params['kappa_fixed']
                    except:
                        d = math.nan
                    class_ = ClassName(d)
            elif noisemodel_name=='Matern':
                try:
                    d = -0.5*control.params['kappa_fixed']
                except:
                    d = math.nan
                ClassName = getattr(sys.modules[__name__], noisemodel_name)
                class_ = ClassName(d)

            else:
                ClassName = getattr(sys.modules[__name__], noisemodel_name)
                class_ = ClassName() 

            self.noisemodels.append(class_)
            if self.verbose==True:
                print('{0:d}) {1:s}'.format(i,noisemodel_name))
            i += 1

        self.Nmodels = len(self.noisemodels)
        self.Nparam = self.Nmodels-1  # weight parameters

        #--- Do we need to estimate additional noise parameters?
        for noisemodel in self.noisemodels:
            self.Nparam += noisemodel.get_Nparam()
        if self.verbose==True:
            print('Nparam : {0:d}'.format(self.Nparam))           


    def get_Nparam(self):
        """ Return the number of parameters to estimate numerically
        
        Returns:
            self.Nparam (int) : total number of parameters 
        """

        return self.Nparam


        
    def compute_fraction(self,i,param):
        """ Compute fraction of noise model i

        Args:
            i (int) : index of noise model
            param (array float) : parameters describing weight of noise models
        
        Returns:
            fraction (float) : fraction of noise model
        """

        #--- Constant
        hpi = 2.0*math.atan(1.0)

        #--- compute fraction
        if self.Nmodels==1:
            return 1.0
        else:
            fraction = 1.0
            for j in range(0,i):
                fraction *= math.sin(hpi*param[j]);
            if i<self.Nmodels-1:
                fraction *= math.cos(hpi*param[i]);

        #--- Some range checks
        if fraction>1.0:
            fraction = 1.0

        return pow(fraction,2.0)
   


    def compute_penalty(self,param):
        """ penalty for each fraction outside the [0:1] range

        Args:
            param (list float) : fractions of each noise model (Nmodels-1)

        Returns:
            penalty value (float)
        """

        #--- Constant
        LARGE = 1.0e8
 
        #--- first fractions
        penalty = 0.0
        for i in range(0,len(self.noisemodels)-1):
            if param[i]<0.0:
                penalty += (0.0-param[i])*LARGE
                param[i] = 0.0
            elif param[i]>1.0:
                penalty += (param[i]-1.0)*LARGE
                param[i] = 1.0

        #--- Extra penalties for noise model parameters
        k = len(self.noisemodels)-1
        for noisemodel in self.noisemodels:
            penalty += noisemodel.penalty(k,param)

        return penalty



    def create_t(self,m,param):
        """

        Args:
            m (int) : length of time series
            param (array float) : array of parameters to estimate

        Returns:
            t (array float) : first row of covariance matrix
        """

        #--- Create empty autocovariance array
        t = np.zeros(m)

        #--- Add autocovariance of each noise model
        k = self.Nmodels-1
        for i in range(0,self.Nmodels):
            fraction = self.compute_fraction(i,param)
            t_part,k_new = self.noisemodels[i].create_t(m,k,param)

            t += fraction * t_part
            k  = k_new

        return t



    def show_results(self,output,noise_params,sigma_eta):
        """ show estimated noiseparameters

        Args:
            output (dictionary) : where values for json file are saved
            noise_params (float-array) : fractions + noise model parameters
            sigma_eta (float) : driving noise value
        """

        if self.verbose==True:
            print('\nNoise Models\n------------')
        k = self.Nmodels-1
        output_block = {}
        for i in range(0,len(self.noisemodels)):
            noisemodel = self.noisemodels[i]
            output_single = {}
            fraction = self.compute_fraction(i,noise_params)

            if self.verbose==True:
                print('{0:s}:'.format(self.noisemodel_names[i]))
                print('fraction  = {0:7.5f}'.format(fraction))

            sigma = math.sqrt(fraction)*sigma_eta
            output_single['fraction'] = fraction
            #--- sigma is stored in the json inside each noise model class
            k = noisemodel.show_results(output_single,k,noise_params,sigma)

            output_block[self.noisemodel_names[i]] = output_single

        output['NoiseModel'] = output_block
