classdef IndexReturns < OptimProblem
  properties 
    name          = 'IndexReturns';
    
    calendar      = [];         % (Optional) calendar associated with the above series
    series        = [];       

    Monetary      = [];         % A placeholder for the (filtered) Monetary object
    
    cond_dist     = 'Normal';   % The (conditional) distribution of the innovations. 
    method        = 'MLE';      % The estimation method

    model_restr   =  1;         % 1 for Wilkie-like, 
                                % 2 for Extended Wilkie-like, 
                                % 3 for ADG-like, 
                                % 4 for full new ESG  
                                
    %%#####  INTERNAL  ############################################################
    mle  = struct();            % A placeholder for MLE results
  end % end properties
    
  methods 
    %% Construct a IndexReturns model 
    function self = IndexReturns(varargin)
     self = self@OptimProblem();
      for no = 1:2:length(varargin)
        setfield(self, varargin{no}, varargin{no+1});
      end      

      % IndexReturns parameters
      self.addParameter('mu_y1',            7.0354e-04,   [-1.000, 1.000]);
      self.addParameter('mu_y2',            0.0067,       [-1.000, 1.000]);
      self.addParameter('mu_y3',            0.0023,       [-1.000, 1.000]);
      self.addParameter('sig2_y1',          0.0020,       [ 1e-9,  0.100]);
      self.addParameter('sig2_y2',          0.0020,       [ 1e-9,  0.100]);
      self.addParameter('sig2_y3',          0.0020,       [ 1e-9,  0.100]);
      self.addParameter('alpha_y',          0.1844,       [ 0.00, 1.0000]);
      self.addParameter('alphabetagamma_y', 0.6745,       [ 0.00, 0.9950]);
      self.addParameter('gamma_y',          1.3405,       [  -10,     10]);
      self.addParameter('sig2_y_init',      0.0055,       [0.0001, 0.100]);
      
      if self.model_restr == 1
        self.params.mu_y2.fixed             = true;
        self.params.mu_y3.fixed             = true;
        self.params.sig2_y2.fixed           = true;
        self.params.sig2_y3.fixed           = true;
        self.params.alpha_y.fixed           = true;
        self.params.alphabetagamma_y.fixed  = true;
        self.params.gamma_y.fixed           = true;
        self.params.sig2_y_init.fixed       = true;

        self.params.alpha_y.value           = 0;
        self.params.alphabetagamma_y.value  = 0;
        self.params.gamma_y.value           = 0;
        self.params.sig2_y_init.value       = self.params.sig2_y1.value;
        self.params.mu_y2.value             = self.params.mu_y1.value;
        self.params.mu_y3.value             = self.params.mu_y1.value;
        self.params.sig2_y2.value           = self.params.sig2_y1.value;
        self.params.sig2_y3.value           = self.params.sig2_y1.value;
      elseif self.model_restr == 2
        self.params.mu_y2.fixed             = true;
        self.params.mu_y3.fixed             = true;
        self.params.sig2_y2.fixed           = true;
        self.params.sig2_y3.fixed           = true;
        self.params.alpha_y.fixed           = true;
        self.params.alphabetagamma_y.fixed  = true;
        self.params.gamma_y.fixed           = true;
        self.params.sig2_y_init.fixed       = true;

        self.params.alpha_y.value           = 0;
        self.params.alphabetagamma_y.value  = 0;
        self.params.gamma_y.value           = 0;
        self.params.sig2_y_init.value       = self.params.sig2_y1.value;
        self.params.mu_y2.value             = self.params.mu_y1.value;
        self.params.mu_y3.value             = self.params.mu_y1.value;
        self.params.sig2_y2.value           = self.params.sig2_y1.value;
        self.params.sig2_y3.value           = self.params.sig2_y1.value;
      elseif self.model_restr == 3
        self.params.alpha_y.fixed           = true;
        self.params.alphabetagamma_y.fixed  = true;
        self.params.gamma_y.fixed           = true;
        self.params.sig2_y_init.fixed       = true;

        self.params.alpha_y.value           = 0;
        self.params.alphabetagamma_y.value  = 0;
        self.params.gamma_y.value           = 0;
        self.params.sig2_y_init.value       = 0;
      elseif self.model_restr == 4
        self.params.sig2_y2.fixed           = true;
        self.params.sig2_y3.fixed           = true;
        
        self.params.sig2_y2.value           = self.params.sig2_y1.value;
        self.params.sig2_y3.value           = self.params.sig2_y1.value;
      end
    end % end IndexReturns
    
    %% This function returns the parameters
    function pv = getPV(self)
      pv         = getPV@OptimProblem(self);

      % Again, depending on the model selected, some parameters are turned
      % on or off.
      if self.model_restr == 1
        self.params.sig2_y_init.value = self.params.sig2_y1.value;
        self.params.mu_y2.value       = self.params.mu_y1.value;
        self.params.mu_y3.value       = self.params.mu_y1.value;
        self.params.sig2_y2.value     = self.params.sig2_y1.value;
        self.params.sig2_y3.value     = self.params.sig2_y1.value;
        
        pv.sig2_y_init                = self.params.sig2_y3.value;
        pv.mu_y2                      = self.params.mu_y1.value;
        pv.mu_y3                      = self.params.mu_y1.value;
        pv.sig2_y2                    = self.params.sig2_y1.value;
        pv.sig2_y3                    = self.params.sig2_y1.value;
      elseif self.model_restr == 2
        self.params.sig2_y_init.value = self.params.sig2_y1.value;
        self.params.mu_y2.value       = self.params.mu_y1.value;
        self.params.mu_y3.value       = self.params.mu_y1.value;
        self.params.sig2_y2.value     = self.params.sig2_y1.value;
        self.params.sig2_y3.value     = self.params.sig2_y1.value;
        
        pv.sig2_y_init                = self.params.sig2_y3.value;
        pv.mu_y2                      = self.params.mu_y1.value;
        pv.mu_y3                      = self.params.mu_y1.value;
        pv.sig2_y2                    = self.params.sig2_y1.value;
        pv.sig2_y3                    = self.params.sig2_y1.value;
      elseif self.model_restr == 4
        self.params.sig2_y2.value     = self.params.sig2_y1.value;
        self.params.sig2_y3.value     = self.params.sig2_y1.value;
        
        pv.sig2_y2                    = self.params.sig2_y1.value;
        pv.sig2_y3                    = self.params.sig2_y1.value;
      end
      
      pv.beta_y = pv.alphabetagamma_y - pv.alpha_y*(1+pv.gamma_y^2);
    end % end getPV

    %% This functions uses fminsearch to find the MLE parameters
    function [results] = fminsearch(self, lambda, varargin)
      results = fminsearch@OptimProblem(self, lambda, varargin);     
      [~,z,sig2] = self.objective(self.getPValues);
      
      switch self.method
        case 'MLE'
          self.mle.params = self.getPV;
          self.mle.z    = z;
          self.mle.sig2 = sig2;
          self.mle.out = results;
      end
    end % end fminsearch

    %% This function computes the objective function; the only case considered so far is the MLE
    function [S,z,sig2] = objective(self, x, varargin)
      switch self.method
        case 'MLE'
          [logl,z,sig2] = self.getMLEfunction(x);
          S = -sum(logl);
      end
    end % end objective

    %% This function computes the likelihood function for the index return model
    function [logl,z,sig2] = getMLEfunction(self,p)
      % We extract the current value of the parameters
      self.setPValues(p);
      pv = self.getPV;
      
      % We extract the number of observations and the monetary policy chain
      T = length(self.series);
      m = [self.Monetary.regimes;3];

      z    = zeros(T,1);
      y    = self.series;
      sig2 = zeros(T+1,1);

      mus    = [pv.mu_y1,pv.mu_y2,pv.mu_y3];
      sig2s  = [pv.sig2_y1,pv.sig2_y2,pv.sig2_y3];
      
      % Change the initial value depending on the model
      switch self.model_restr
        case 1
          sig2(1) = sig2s(1);
        case 2
          sig2(1) = sig2s(1);
        case 3
          sig2(1) = sig2s(floor(self.Monetary.params.m0.value));
        case 4
          sig2(1) = pv.sig2_y_init;
        otherwise
          error('Model has not been implemented yet.');
      end

      % For each observation, we compute the log-likelihood contribution at
      % time t
      for dt = 1:T
        z(dt) = y(dt) - mus(m(dt)) + 0.5.*sig2(dt);
        sig2(dt+1) = sig2s(m(dt)) + pv.beta_y*(sig2(dt)-sig2s(m(dt))) + pv.alpha_y*( (z(dt) - pv.gamma_y.*sqrt(sig2(dt))).^2 - (1+pv.gamma_y.^2)*sig2s(m(dt)));
      end
      sig2 = sig2(1:end-1);
      z    = z(1:end);
    	logl = log( normpdf( z, 0, sqrt(sig2)) );
      
      % Hack: if a loglikelihood is not real, we return an array of NaN
      if imag(sum(logl)) ~= 0
        logl = NaN(size(logl));
      end
    end % end getMLEfunction

    %% This function computes the log prior for the index return model
    function logprior = getlogPrior(self,pv)      
      logprior =  log(normpdf(pv.mu_y1,0,1)) + ...
                  log(normpdf(pv.mu_y2,0,1)) + ...
                  log(normpdf(pv.mu_y3,0,1)) + ...
                  log(exppdf(pv.sig2_y1,1)) + ...
                  log(exppdf(pv.sig2_y2,1)) + ...
                  log(exppdf(pv.sig2_y3,1)) + ...
                  log((pv.alpha_y <= 0.995).*(pv.alpha_y >= 0)) + ...               		  % Hack: Otherwise, we might have some stationarity issues...
                  log((pv.alphabetagamma_y <= 0.995).*(pv.alphabetagamma_y >= 0)) + ...       % Hack: Otherwise, we might have some stationarity issues...
                  log((pv.beta_y <= 0.995).*(pv.beta_y >= 0)) + ...    						  % Hack: Otherwise, we might have some stationarity issues...
                  log(normpdf(pv.gamma_y,0,10)) + ...
                  log(exppdf(pv.sig2_y_init,1));
    end % end getlogPrior    
  end % end methods 
end % end IndexReturns
