#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed May 15 10:07:33 2024

@author: jdaniel
"""

import numpy as np
import pandas as pd

def get_consumption():
    iscale = 0 #if 1 then scale consumption to production, otherwise no scaling of consumption
    iregions = 10
    files = 'productionConsumption.xlsx'
    sheets = ['Total production','Import-export','Feedstock']
    ilen = 34
    inputs = np.zeros([len(sheets),ilen,iregions])
    years = np.zeros(ilen)
    for i in range(0,len(sheets)):
        fileinfo = pd.read_excel(files,sheet_name=sheets[i])
        fileinfo.fillna(0, inplace=True)
        inputs[i,:,:] = fileinfo.iloc[:,1::]
        years[:] = fileinfo.iloc[:,0]
    regions = fileinfo.columns[1::]
    production = inputs[0,:,:] - inputs[2,:,:]
    production[1,2::] = production[0,2::]*2/3 + production[3,2::]*1/3
    production[2,2::] = production[0,2::]*1/3 + production[3,2::]*2/3
    consumption = production + inputs[1,:,:] 
    feedstock = inputs[2,:,:]
    p_sum = np.sum(production,axis=1)
    c_sum = np.sum(consumption,axis=1)
    for i in range(0,len(p_sum)):
        if (iscale): 
            consumption[i,:] = consumption[i,:]*p_sum[i]/c_sum[i]
    adddata = 1
    if (adddata == 1):
        years = np.append(years,years[-1]+[1,2,3,4,5])  
        cshape = np.shape(consumption)  
        cnew = np.zeros([cshape[0]+5,cshape[1]])    
        cnew[0:cshape[0],:] = consumption
        tmp = np.zeros(ilen-29)
        fit = np.polyfit(years[26:34],np.sum(cnew[26:34,:],axis=1),1)
        tmp[:] = fit[0]*years[34::] + fit[1]
        ii = np.where(tmp >= 0)
        print(ii)
        cnew[34+ii[0],7] = tmp[ii]
        for i in range(0,-1):
            fit = np.polyfit(years[26:34],cnew[26:34,i],1)
            tmp[:] = fit[0]*years[34::] + fit[1]
            print(i,fit)
            print(i,cnew[26:34,i],tmp)
            print(i,consumption[26:34,i])
            ii = np.where(tmp >= 0)
            print(tmp[ii])
            if (len(ii[0] > 0)):
                if (fit[0] < 0):
                    cnew[34+ii[0],i] = tmp[ii]
    else:
        cnew = consumption
    return(years,production,cnew,feedstock,regions)

def getSegmentedVals(years,consumption,marketSegRefrig,marketSegNoRefrig):
    arrsize = consumption.shape
    arrsize2 = marketSegRefrig.shape
    c_out = np.zeros([arrsize[0],arrsize[1],arrsize2[1]])
    c_use = np.zeros([arrsize[0],arrsize[1],arrsize2[1]])
    c_outNoRefrig = np.zeros([arrsize[0],arrsize[1],arrsize2[1]])
    for i in range(0,arrsize[0]):
        for j in range(0,arrsize[1]):
            c_out[i,j,:] = consumption[i,j] * marketSegRefrig[j,:]
            c_outNoRefrig[i,j,:] = consumption[i,j] * marketSegNoRefrig[j,:]
    for i in range(0,arrsize[0]):
        if (years[i] <= 2010):
            c_use[i,:,:] = c_out[i,:,:]
        elif (years[i] <=2015):
            c_use[i,:,:] = c_out[i,:,:] + (years[i]-2010)/5. * (c_outNoRefrig[i,:,:] - c_out[i,:,:])
        else:
            c_use[i,:,:] = c_outNoRefrig[i,:,:]
    return(c_use)

def getSegmentedValsMult(years,consumption,ms1,ms2,ms3,msno1,msno2,msno3):
    arrsize = consumption.shape
    arrsize2 = ms1.shape
    c_out = np.zeros([arrsize[0],arrsize[1],arrsize2[1]])
    c_use = np.zeros([arrsize[0],arrsize[1],arrsize2[1]])
    c_outNoRefrig = np.zeros([arrsize[0],arrsize[1],arrsize2[1]])
    for i in range(0,arrsize[0]):
        for j in range(0,arrsize[1]):
            if (years[i] <= 2002):
                c_out[i,j,:] = consumption[i,j]*ms1[j,:]
                c_outNoRefrig[i,j,:] = consumption[i,j]*msno1[j,:]
            elif (years[i] <= 2007):
                c_out[i,j,:] = consumption[i,j]*ms2[j,:]
                c_outNoRefrig[i,j,:] = consumption[i,j]*msno2[j,:]
            else:
                c_out[i,j,:] = consumption[i,j]*ms3[j,:]
                c_outNoRefrig[i,j,:] = consumption[i,j]*msno3[j,:]
    for i in range(0,arrsize[0]):
        if (years[i] <= 2010):
            c_use[i,:,:] = c_out[i,:,:]
        elif (years[i] <=2015):
            c_use[i,:,:] = c_out[i,:,:] + (years[i]-2010)/5. * (c_outNoRefrig[i,:,:] - c_out[i,:,:])
        else:
            c_use[i,:,:] = c_outNoRefrig[i,:,:]
    return(c_use)