# Python 2.7 code
# 
#
# ----------------------------- Python Configuration --------------------------
from __future__ import division
import numpy as np
import time
import pylab
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import os
from statsmodels.tsa.filters import hpfilter as hp
import scipy.stats
from scipy import random
from scipy.stats.stats import pearsonr
from scipy.stats import scoreatpercentile as perc

# ----------------------------- Preliminaries --------------------------
### Load up the data
# Monthly
employment_raw_raw = np.genfromtxt('E_tot_2007.csv',delimiter=',')
T = np.size(employment_raw_raw)
employment_raw = employment_raw_raw[:T]
unemployment_short_raw = np.genfromtxt('U_shortterm_2007.csv', delimiter=',')[:T]
unemployment_long_raw = np.genfromtxt('U_longterm_2007.csv', delimiter=',')[:T]
nonparticipation_raw = np.genfromtxt('N_tot_2007.csv', delimiter=',')[:T]
psi_employment = np.genfromtxt('psi_e_2007.csv', delimiter=',')[:T]
psi_unemployment_short = np.genfromtxt('psi_ushort_2007.csv', delimiter=',')[:T]
psi_unemployment_long = np.genfromtxt('psi_ulong_2007.csv', delimiter=',')[:T]
psi_nonparticipation = np.genfromtxt('psi_n_2007.csv', delimiter=',')[:T]
unemployment_0_raw = np.genfromtxt('U(0)_2007.csv',delimiter=',')[:T] / 1000
unemployment_6_raw = np.genfromtxt('U(6)_2007.csv',delimiter=',')[:T] / 1000
unemployment_7_raw = np.genfromtxt('7months_2007.csv',delimiter=',')[:T] / 1000
# Annual
theta_eu_0_annual = np.genfromtxt('theta_eu(0)_2007.csv', delimiter=',') / 100
theta_eu_6_annual = np.genfromtxt('theta_eu(6)_2007.csv', delimiter=',') / 100
theta_nu_0_annual = np.genfromtxt('theta_nu(0)_2007.csv', delimiter=',') / 100
theta_nu_6_annual = np.genfromtxt('theta_nu(6)_2007.csv', delimiter=',') / 100

### Indexing
T_annual = np.size(theta_eu_0_annual)

### Time aggregate theta
# Preallocate
theta_eu_0, theta_eu_6, theta_nu_0, theta_nu_6 = np.zeros(T), np.zeros(T), np.zeros(T), np.zeros(T)
# Initial year (starts March 2002)
theta_eu_0[:12], theta_eu_6[:12], theta_nu_0[:12], theta_nu_6[:12] = theta_eu_0_annual[0], theta_eu_6_annual[0], theta_nu_0_annual[0], theta_nu_6_annual[0]
# Loop
for t in range(T_annual-1):
    theta_eu_0[12+12*t:12+12*(t+1)], theta_eu_6[12+12*t:12+12*(t+1)], theta_nu_0[12+12*t:12+12*(t+1)], theta_nu_6[912+12*t:12+12*(t+1)] = theta_eu_0_annual[t+1], theta_eu_6_annual[t+1], theta_nu_0_annual[t+1], theta_nu_6_annual[t+1]

### Normalize population size
population = employment_raw + unemployment_short_raw + unemployment_long_raw + nonparticipation_raw
employment, unemployment_short, unemployment_long, unemployment_0, unemployment_6, unemployment_7, nonparticipation = employment_raw / population, unemployment_short_raw / population, unemployment_long_raw / population, unemployment_0_raw / population, unemployment_6_raw / population, unemployment_7_raw / population, nonparticipation_raw / population

### Create net flows
delta_employment = employment[1:] - employment[:T-1]
delta_nonparticipation = nonparticipation[1:] - nonparticipation[:T-1]

# ----------------------------- Compute the lambdas --------------------------
### Preallocate
lambda_vector = np.zeros((8,T-1))
### Do the loop
for t in range(T-1):
    # Preallocate matrices
    A_matrix, B_matrix = np.zeros((8,8)), np.zeros((8,1))
    # Input entries by hand: A_matrix
    A_matrix[0,0], A_matrix[0,2], A_matrix[0,4], A_matrix[0,6], A_matrix[0,7] = employment[t], unemployment_short[t], unemployment_long[t], -nonparticipation[t], -nonparticipation[t]
    A_matrix[1,0], A_matrix[1,1], A_matrix[1,3], A_matrix[1,5], A_matrix[1,7] = -employment[t], -employment[t], unemployment_short[t], unemployment_long[t], nonparticipation[t]
    A_matrix[2,1], A_matrix[2,6] = theta_eu_0[t] * employment[t], theta_nu_0[t] * nonparticipation[t]
    A_matrix[3,1], A_matrix[3,4], A_matrix[3,5], A_matrix[3,6] = theta_eu_6[t] * employment[t], -unemployment_6[t], -unemployment_6[t], theta_nu_6[t] * nonparticipation[t]
    A_matrix[4,6], A_matrix[4,7] = 1, -psi_nonparticipation[t]
    A_matrix[5,4], A_matrix[5,5] = 1, -psi_unemployment_long[t]
    A_matrix[6,2], A_matrix[6,3] = 1, -psi_unemployment_short[t]
    A_matrix[7,0], A_matrix[7,1] = 1, -psi_employment[t]
    # Input entries by hand: B_matrix
    B_matrix[0], B_matrix[1], B_matrix[2], B_matrix[3] = delta_nonparticipation[t], delta_employment[t], unemployment_0[t+1], unemployment_7[t+1] - unemployment_6[t]
    # Solve the equation
    lambda_vector[:,t] = np.dot(np.linalg.inv(A_matrix),B_matrix).flatten()

# ----------------------------- Store the output --------------------------
np.savetxt('lambda_en_output_2007.csv',lambda_vector[0,:],delimiter=',')
np.savetxt('lambda_eu_output_2007.csv',lambda_vector[1,:],delimiter=',')
np.savetxt('lambda_usn_output_2007.csv',lambda_vector[2,:],delimiter=',')
np.savetxt('lambda_use_output_2007.csv',lambda_vector[3,:],delimiter=',')
np.savetxt('lambda_uln_output_2007.csv',lambda_vector[4,:],delimiter=',')
np.savetxt('lambda_ule_output_2007.csv',lambda_vector[5,:],delimiter=',')
np.savetxt('lambda_nu_output_2007.csv',lambda_vector[6,:],delimiter=',')
np.savetxt('lambda_ne_output_2007.csv',lambda_vector[7,:],delimiter=',')
"""
# ----------------------------- Transition back to average -----------------
### Time it takes
cap_T = 12 * 3

### Compute averages
lambda_en_avg, lambda_eu_avg, lambda_un_avg, lambda_ue_avg, lambda_nu_avg, lambda_ne_avg = np.mean(lambda_vector[0,:71]), np.mean(lambda_vector[1,:71]), np.mean(lambda_vector[2,:71]), np.mean(lambda_vector[3,:71]), np.mean(lambda_vector[4,:71]), np.mean(lambda_vector[5,:71])

### Transition back to average
t_range = np.linspace(0,cap_T,cap_T)
lambda_en_transition = (lambda_vector[0,T-3] * (cap_T - t_range) + lambda_en_avg * t_range) / cap_T
lambda_eu_transition = (lambda_vector[1,T-3] * (cap_T - t_range) + lambda_eu_avg * t_range) / cap_T
lambda_un_transition = (lambda_vector[2,T-3] * (cap_T - t_range) + lambda_un_avg * t_range) / cap_T
lambda_ue_transition = (lambda_vector[3,T-3] * (cap_T - t_range) + lambda_ue_avg * t_range) / cap_T
lambda_nu_transition = (lambda_vector[4,T-3] * (cap_T - t_range) + lambda_nu_avg * t_range) / cap_T
lambda_ne_transition = (lambda_vector[5,T-3] * (cap_T - t_range) + lambda_ne_avg * t_range) / cap_T

### Save output
np.savetxt('lambda_en_transition.csv',lambda_en_transition,delimiter=',')
np.savetxt('lambda_eu_transition.csv',lambda_eu_transition,delimiter=',')
np.savetxt('lambda_un_transition.csv',lambda_un_transition,delimiter=',')
np.savetxt('lambda_ue_transition.csv',lambda_ue_transition,delimiter=',')
np.savetxt('lambda_nu_transition.csv',lambda_nu_transition,delimiter=',')
np.savetxt('lambda_ne_transition.csv',lambda_ne_transition,delimiter=',')
"""
