#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Oct 22 11:10:33 2021

@author: rishabh
"""

import pandas as pd
import pickle
import numpy as np

# count by varlist
def get_unique(dset, varlist, dropna = 0):
    if dropna:
        temp = dset.dropna()
    else:
        temp = dset
    return temp.groupby(varlist).size().reset_index().rename(columns={0:'count'})

def movecol(df, cols_to_move=[], ref_col='', place='After'):

    cols = df.columns.tolist()
    if place == 'After':
        seg1 = cols[:list(cols).index(ref_col) + 1]
        seg2 = cols_to_move
    if place == 'Before':
        seg1 = cols[:list(cols).index(ref_col)]
        seg2 = cols_to_move + [ref_col]

    seg1 = [i for i in seg1 if i not in seg2]
    seg3 = [i for i in cols if i not in seg1 + seg2]

    return(df[seg1 + seg2 + seg3])

def split_str(seq, chunk, skip_tail=False):
    lst = []
    if chunk <= len(seq):
        lst.extend([seq[:chunk]])
        lst.extend(split_str(seq[chunk:], chunk, skip_tail))
    elif not skip_tail and seq:
        lst.extend([seq])
    return lst

# Gives daterange for which values are present for list of countries and series in df in tidy_format
def get_daterange(df):

    countries = get_unique(df, ['Country'])
    series = get_unique(df, ['SeriesCode'])
    date_range  = pd.DataFrame(index=countries['Country'], columns = series['SeriesCode'])

    for s, ser in enumerate(series['SeriesCode'].to_list()):
        for c, cou in enumerate(countries['Country'].to_list()):
            try:
                temp = df.loc[(df.Country== cou) & (df.SeriesCode == ser)]
                date1 = temp.loc[temp['value'].first_valid_index(), 'Date'].to_period('Q').strftime('%FQ%q')
                date2 = temp.loc[temp['value'].last_valid_index(), 'Date'].to_period('Q').strftime('%FQ%q')
                date_range.loc[cou, ser] = [date1+'-'+date2]
            except:
                date_range.loc[cou, ser] = 'None'

    return date_range

def save_pickle(df, filename):
    temp = {}
    temp['name'] = df
    with open(filename, 'wb') as handle:
        pickle.dump(temp, handle, protocol=pickle.HIGHEST_PROTOCOL)

def load_pickle(filename):
    with open(filename, 'rb') as handle:
        #temp = pickle.load(handle)
        temp=pd.read_pickle(handle)
    return temp['name']

def load_pickle_old(filename):
    with open(filename, 'rb') as handle:
        temp = pickle.load(handle)
    return temp['name']


# Print and data check functions
def print_ss(d):
    return pd.DataFrame(d.items()).set_index(0).rename(columns={1:'value'}).T

def print_aggregates(d, extra_vars = [], extra_only = 0):
    aggregate_vars = ['realrate', 'r', 'B', 'G', 'Y', 'C', 'Cstar', 'Q', 'nfa', 'A', 'T', 'Z', 'Transfer', 'w', 'E', 'pi', 'rstar', 'rss', 'i', 'piHH']
    if extra_only == 0:
        aggregate_vars = sorted(list(set(aggregate_vars) | set(extra_vars)))
    else:
        aggregate_vars = extra_vars
    dict_you_want={}
    for key in aggregate_vars:
        if key in d:
            dict_you_want[key] = d[key]
    return pd.DataFrame(dict_you_want.items()).set_index(0).rename(columns={1:'value'}).T
 

def compute_world_sum(ss, varlist, weight_var = [], mul=1):
    
    """Computes the weighted sum of varlist in all ss countries"""
    
    which_countries = list(ss.keys())
    aggregates = np.zeros(len(varlist))
    
    for v, vari in enumerate(varlist):
        for c, cou in enumerate(which_countries) :
            
            if weight_var == []:
                aggregates[v]+=ss[cou][vari]
            else:
                if mul==1:
                    aggregates[v]+=ss[cou][vari]*ss[cou][weight_var[v]]
                else:
                    aggregates[v]+=ss[cou][vari]/ss[cou][weight_var[v]]
    return aggregates
        
def print_var_by_c(ss, varlist):
    
    df = {}
    which_countries = list(ss.keys())
    df['Code'] = which_countries 
    for v, vari in enumerate(varlist):
        temp = np.zeros(len(which_countries))
        for c, cou in enumerate(which_countries) :
            temp[c] = ss[cou][vari]
    
        df[vari] = temp
        del temp
     
    return pd.DataFrame(df).set_index('Code')


