import numpy as np
import pandas as pd
import gen_fns
import math
import re
import csv
import correlation_matrix as co


#
# reformat_raw_data
#
# Takes compiled data on historical proliferation contributers from xlsx
# spreadsheet (final_hist_data.xlsx) and reformats. xlsx file has a single
# row for each country with separate data for a 'pursue date' and an 'acquire
# 'date'.  These are split out so that each state has a unique row entry for
# every year for which we have data. Nuclear Weapon states have an entry when
# they began pursuing and another when they acquired. States that explored or
# pursued have an entry for that year as well as an entry for 2015. States
# that never pursued have only an entry for 2015.
# A 'Status' column in output codes the action of the state in the given year
# (Never/Explore/Pursue/Acquire). Coding details available inline.
#
# Input
#    File:       tab-separated-file dropped straight from xlsx spreadsheet
#    (n_header): number of extra header lines to be skipped from spreadsheet
#    (outfile):  name of desired output file, written to directory you're
#                running python from
# Output
#    States:     Array of states for each entry of the file database
#    Data:       structured array of historical data for each state at time
#                indicated by 'Date' 
#
def reformat_raw_data(file, n_header=1, outfile=None):
    from gen_fns import get_data
    from numpy.lib.recfunctions import append_fields
    
    countries, columns, raw_data = get_data(file, n_header, named_struct=True)

    pursue_names = []
    acquire_names = []
    clean_names = []

    status = np.full(len(countries), 0)
    raw_data = append_fields(raw_data, 'Status', status)

    cols = raw_data.dtype.names
    order = range(0, len(cols) - 1)
    order.insert(2,len(cols) - 1)
    data = raw_data[[cols[i] for i in order]]
    cols = data.dtype.names
    
    for c in range(len(cols)):
        if ('Pursuit_' in cols[c]):
            pursue_names.append(cols[c])
            new_str = re.sub('Pursuit_', '', cols[c])
            clean_names.append(new_str)
        elif ('Acquire_' in cols[c]):
            acquire_names.append(cols[c])
        else:
            pursue_names.append(cols[c])
            acquire_names.append(cols[c])
            clean_names.append(cols[c])
        
    # all countries, data split into pursue-relevant or acquire-relevant
    pursue_array = data[pursue_names]
    acquire_array = data[acquire_names]

    no_acquire_mask = np.isnan(acquire_array['Acquire_Date'])
    conven_mask = np.isnan(pursue_array['Pursuit_Date'])
    explore_mask = (pursue_array['Pursuit_Date'] < 0)
    pursue_only_mask = (np.isnan(acquire_array['Acquire_Date'])) & (~(np.isnan(pursue_array['Pursuit_Date'])) & (pursue_array['Pursuit_D
ate'] > 0))

    # states that pursued (_init_) and successfully acquired
    acquire_states = countries[~no_acquire_mask]
    acquire_final_array = acquire_array[~no_acquire_mask]
    acquire_init_array = pursue_array[~no_acquire_mask]

    # states that never pursued and have only conventional weapons
    conven_array = pursue_array[conven_mask]
    conven_states = countries[conven_mask]

    
# states that explored NW but ultimately did not pursue
    explore_array = pursue_array[explore_mask]
    explore_present_array = acquire_array[explore_mask]
    explore_states = countries[explore_mask]

    # states that pursued but did not succeeed in acquiring
    pursue_only_array = pursue_array[pursue_only_mask]
    pursue_present_array = acquire_array[pursue_only_mask]
    pursue_states = countries[pursue_only_mask]

    # Status
    # -1 present data for a state that unsucessfully pursued
    # 0 present data for never-pursued (has only 'conventional weapons')
    # 1 historical data for explored
    # 2 historical data for pursued
    # 3 historical data for acquired
    acquire_init_array['Status'] = 2
    acquire_final_array['Status'] = 3
    conven_array['Status'] = 0
    pursue_only_array['Status'] = 2
    pursue_present_array['Status'] = -1
    explore_array['Status'] = 1
    explore_present_array['Status'] = 0
    
    conven_array['Pursuit_Date'] = 2015
    explore_array['Pursuit_Date'] = abs(explore_array['Pursuit_Date'])
    explore_present_array['Acquire_Date'] = 2015
    pursue_present_array['Acquire_Date'] = 2015
    
    acquire_final_array.dtype.names = clean_names
    acquire_final_array.mask.dtype.names = clean_names
    acquire_init_array.dtype.names = clean_names
    acquire_init_array.mask.dtype.names = clean_names
    conven_array.dtype.names = clean_names
    conven_array.mask.dtype.names = clean_names
    pursue_only_array.dtype.names = clean_names
    pursue_only_array.mask.dtype.names = clean_names
    pursue_present_array.dtype.names = clean_names
    pursue_present_array.mask.dtype.names = clean_names
    explore_array.dtype.names = clean_names
    explore_array.mask.dtype.names = clean_names
    explore_present_array.dtype.names = clean_names
    explore_present_array.mask.dtype.names = clean_names

    
    final_states = np.hstack((conven_states,
                              explore_states,
                              explore_states,
                              pursue_states,
                              acquire_states,
                              acquire_states,
                              pursue_states))
    final_data = np.hstack((conven_array,
                            explore_present_array,
                            explore_array,
                            pursue_only_array,
                            acquire_init_array,
                            acquire_final_array,
                            pursue_present_array))
    
    header ='Country' + '\t' + ('\t'.join(map(str,final_data.dtype.names)))

    if (outfile != None):
        with open(outfile, 'wb') as f:
            writer = csv.writer(f)
            writer.writerow([header])
            i = 0
            for comp in final_data.compressed():
                cur_line = final_states[i]
                for c in range(len(comp)):
                    val = comp[c]
                    if (c <= 1):
                        val = int(val)
                    cur_line = cur_line + '\t' + str(val)
                writer.writerow([cur_line])
                i+=1
                
    return final_states,final_data


#
# calc_weights
#
# Given a tab separated input file with the following columns, use partial
# component analysis to determine the relative importance of each factor
# In:
#    filename:   tab-separated input file with factor scores for each state
#                columns: Country, Date, Status, Factor1, Factor2, ....
#    mn_status:  min value of nuclear weapons status for state to be considered
#                in analysis (-1 = gave up weapons program, 0 = never pursued,
#                1 = explored, 2 = pursued, 3 = acquired)
#    mx_status:  max value of weapons status to be considered
#    correl_min: weights below this number are rescaled to zero as they are
#                not significant (including negative correlations)
#
# Out:
#    weights:    relative importance of each factor from the input file. Weights
#                sum to 1 and cannot be negative.
#
def calc_weights(filename, mn_status=0, mx_status=2, correl_min = 1e-6):
    data_file = open(filename, 'r')
    full_matrix = np.loadtxt(data_file, skiprows=1,usecols=(2,3,4,5,6,7,8,9,10))
    relevant_mask = ((full_matrix[:,0] >= mn_status) &
                     (full_matrix[:,0] <= mx_status))
    matrix = full_matrix[relevant_mask]

    cor = co.Cor_matrix(matrix)
    factor_vals = np.array(cor[0,1:])[0]
    factor_vals[factor_vals < correl_min] = 0
    f_tot = factor_vals.sum()
    weights = factor_vals/f_tot  # normalize weights to sum to one

    return weights


    
#
# calc_pursuit
#
# Calculates 0-10 scale pursuit scores for historical data using factor
# weighting from Baptiste's correlation analysis.
#
# In:
#     raw_data:   np.array of value for each factor on a 0-10 scale
#     weights:    np.array of factor weights (from correlation analysis)
# Out:
#     final_vals: List of pursuit eq