# Python 2.7 code written by Thomas Winberry
# This draft: February 10th
#
# ----------------------------- 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
# Quarterly
employment_raw_raw = np.genfromtxt('E_tot_to_tom_2013.csv',delimiter=',')
T = np.size(employment_raw_raw)
employment_raw = employment_raw_raw[:T]
unemployment_short_raw = np.genfromtxt('U_shorterm_tom_2013.csv', delimiter=',')[:T]
unemployment_long_raw = np.genfromtxt('U_longterm_tom_2013.csv', delimiter=',')[:T]
nonparticipation_raw = np.genfromtxt('N_tot_to_tom_2013.csv', delimiter=',')[:T]
psi_employment = np.genfromtxt('psi_e.csv', delimiter=',')[:T]
psi_unemployment_short = np.genfromtxt('psi_ushort.csv', delimiter=',')[:T]
psi_unemployment_long = np.genfromtxt('psi_ulong.csv', delimiter=',')[:T]
psi_nonparticipation = np.genfromtxt('psi_n.csv', delimiter=',')[:T]
unemployment_0_raw = np.genfromtxt('U(0)_to_tom_2013_new.csv',delimiter=',')[:T] / 1000
unemployment_6_raw = np.genfromtxt('U(6)_to_tom_2013_new.csv',delimiter=',')[:T] / 1000
unemployment_7_raw = np.genfromtxt('7_month_to_tom_2013_new.csv',delimiter=',')[:T] / 1000
theta_eu_0 = np.genfromtxt('theta_eu_0_monthly.csv',delimiter=',')
theta_eu_7 = np.genfromtxt('theta_eu_7_monthly.csv',delimiter=',')
theta_nu_0 = np.genfromtxt('theta_nu_0_monthly.csv',delimiter=',')
theta_nu_7 = np.genfromtxt('theta_nu_7_monthly.csv',delimiter=',')

### 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_7[t] * employment[t], -unemployment_6[t], -unemployment_6[t], theta_nu_7[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.csv',lambda_vector[0,:],delimiter=',')
np.savetxt('lambda_eu_output.csv',lambda_vector[1,:],delimiter=',')
np.savetxt('lambda_usn_output.csv',lambda_vector[2,:],delimiter=',')
np.savetxt('lambda_use_output.csv',lambda_vector[3,:],delimiter=',')
np.savetxt('lambda_uln_output.csv',lambda_vector[4,:],delimiter=',')
np.savetxt('lambda_ule_output.csv',lambda_vector[5,:],delimiter=',')
np.savetxt('lambda_nu_output.csv',lambda_vector[6,:],delimiter=',')
np.savetxt('lambda_ne_output.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_usn_avg, lambda_use_avg, lambda_uln_avg, lambda_ule_avg, lambda_nu_avg, lambda_ne_avg = np.mean(lambda_vector[0,:9*12]), np.mean(lambda_vector[1,:9*12]), np.mean(lambda_vector[2,:9*12]), np.mean(lambda_vector[3,:9*12]), np.mean(lambda_vector[4,:9*12]), np.mean(lambda_vector[5,:9*12]), np.mean(lambda_vector[6,:9*12]), np.mean(lambda_vector[7,:9*12])

### 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_usn_transition = (lambda_vector[2,T-3] * (cap_T - t_range) + lambda_usn_avg * t_range) / cap_T
lambda_use_transition = (lambda_vector[3,T-3] * (cap_T - t_range) + lambda_use_avg * t_range) / cap_T
lambda_uln_transition = (lambda_vector[4,T-3] * (cap_T - t_range) + lambda_uln_avg * t_range) / cap_T
lambda_ule_transition = (lambda_vector[5,T-3] * (cap_T - t_range) + lambda_ule_avg * t_range) / cap_T
lambda_nu_transition = (lambda_vector[6,T-3] * (cap_T - t_range) + lambda_nu_avg * t_range) / cap_T
lambda_ne_transition = (lambda_vector[7,T-3] * (cap_T - t_range) + lambda_ne_avg * t_range) / cap_T

### Save output
np.savetxt('lambda_en_transition_2010.csv',lambda_en_transition,delimiter=',')
np.savetxt('lambda_eu_transition_2010.csv',lambda_eu_transition,delimiter=',')
np.savetxt('lambda_usn_transition_2010.csv',lambda_usn_transition,delimiter=',')
np.savetxt('lambda_use_transition_2010.csv',lambda_use_transition,delimiter=',')
np.savetxt('lambda_uln_transition_2010.csv',lambda_uln_transition,delimiter=',')
np.savetxt('lambda_ule_transition_2010.csv',lambda_ule_transition,delimiter=',')
np.savetxt('lambda_nu_transition_2010.csv',lambda_nu_transition,delimiter=',')
np.savetxt('lambda_ne_transition_2010.csv',lambda_ne_transition,delimiter=',')
"""
