# find bond-graph parameters whole whole cell with multiple modules
# present:
#  ['cAMP','LRGbinding_B1AR','LRGbinding_M2','GsProtein','GiProtein']
#
# based on Pan cardiac AP 2018 code
#
# find kf and kr through a single file that accounts for potential closed
# loops in the new composite model

import os
import sys
import csv
import math
import numpy as np
from scipy.linalg import null_space
import sympy
from grand_kinetic_parameters import grand_kinetic_parameters
# from grand_kinetic_parameters_Gs import grand_kinetic_parameters

def read_IDs(path):
    data = []
    with open(path,'r') as f:
        reader = csv.reader(f)
        for row in reader:
            data.append(row[0])
        f.close()
    return data

def load_matrix(stoich_path):
    matrix = []
    with open(stoich_path,'r') as f:
        reader = csv.reader(f,delimiter=',')
        for row in reader:
            matrix.append([int(r) for r in row])
        f.close()
    return matrix


def calcT(I_vec,num_rows):
    num_cols = len(I_vec)
    T = np.zeros([num_rows,num_cols])
    for i in range(num_cols):
        T[I_vec[i]][i] = 1
    
    return T


if __name__ == "__main__":
    # Set directories
    current_dir = os.getcwd()
    main_dir = os.path.dirname(current_dir)
    output_dir = current_dir + '\output'

    # Define constants
    R = 8.314 # unit J/mol/K
    T = 310
    F = 96485
    N_A = 6.022e23

    combined_kinetic_parameters = True # all listed modules must be present

    if not combined_kinetic_parameters:
        print("imports of same file name is not correct")

    # Define volumes (unit pL)
    V_myo = 34.4
    V_e = 5.182     # external volume
    V_SR = V_myo*0.035 # SR volume
    V_di = V_myo*0.0539 # diadic volume
    V = dict()
    V['V_myo'] = V_myo
    V['V_SR'] = V_SR
    V['V_di'] = V_di
    V['V_e'] = V_e

    # Load stoichiometric matrices and kinetic rate constants
    array_names = [ 'LRGbinding_M2', 'GiProtein'] # GProtein_forPeter

    num_subsystems = len(array_names)
    sys_struct = {c:{} for c in array_names}
    rxnIDs = []
    Knames = []
    Kname_modules = dict()
    for i_system in range(num_subsystems):
        sys_name = array_names[i_system]
        # sys.path.insert(1,main_dir + '\\' +sys_name+ '\\parameter_finder')
        os.chdir(main_dir + '\\' +sys_name+ '\\parameter_finder')
        forward_mat_path = 'data\\all_forward_matrix.txt'
        reverse_mat_path = 'data\\all_reverse_matrix.txt'
        N_f = load_matrix(forward_mat_path)
        N_r = load_matrix(reverse_mat_path)
        sys_struct[sys_name]['N_f'] = N_f
        sys_struct[sys_name]['N_r'] = N_r

    #     print(array_names[i_system])
        dims = dict()
        dims['num_rows'] = len(N_f)
        dims['num_cols'] = len(N_f[0])
        I = np.identity(dims['num_cols'])
        M = [I, np.transpose(N_f), I, np.transpose(N_r)]
        if not combined_kinetic_parameters:
            from kinetic_parameters import kinetic_parameters
            [k_kinetic, N_cT, K_C, W] = kinetic_parameters(M, True, dims, V)
            sys_struct[sys_name]['kfkr'] = k_kinetic
            sys_struct[sys_name]['Kc'] = K_C
            sys_struct[sys_name]['N_cT'] = N_cT

        rxnID = read_IDs('data\\rxnID.txt')
        rxnIDs.extend(rxnID)
        sys_struct[sys_name]['rxnID'] = rxnID
        Kname = read_IDs('data\\Kname.txt')
        Knames.extend(Kname)
        sys_struct[sys_name]['Kname'] = Kname

        if not combined_kinetic_parameters:
            pass
            # sys.path.remove(current_dir + '\\' +sys_name+ '\\parameter_finder')


    Kunique = []
    for ik in Knames:
        # if ~any(strcmp(Kunique,ik)):
        if ik not in Kunique:
            Kunique.append(ik)

    os.chdir(current_dir)

    # relations between submodule to whole module
    # [['cAMP','LRGbinding_B1AR','LRGbinding_M2','GsProtein','GiProtein']]
    for name in array_names:
        ids = [Kunique.index(kid) for kid in sys_struct[name]['Kname']]
        sys_struct[name]['I_vec'] = ids

    # sys_struct[1].I_vec = 1:18           # cAMP
    # sys_struct[2].I_vec = 19:24          # LRGbinding_B1AR
    # sys_struct[3].I_vec = 25:30          # LRGbinding_M2
    # sys_struct[4].I_vec = [20 21 23 22 24 14 31:33]     # GsProtein
    # sys_struct[5].I_vec = [26 27 29 28 30 16 34 35 33]  # GiProtein

    num_rows = max(sys_struct[array_names[-1]]['I_vec'])+1

    N_f = []
    N_r = []

    for sys_name in array_names:
        # print(sys_name)
        T = calcT(sys_struct[sys_name]['I_vec'],num_rows)
        sys_struct[sys_name]['T'] = T

        # N_f = [N_f, T*sys_struct[sys_name]['N_f']]
        # N_r = [N_r, T*sys_struct[sys_name]['N_r']]
        new_f = np.matmul(T,sys_struct[sys_name]['N_f'])
        new_r = np.matmul(T,sys_struct[sys_name]['N_r'])

        if not len(N_f):
            N_f = new_f
            N_r = new_r
        else:
            N_f = np.append(N_f, new_f,1)
            N_r = np.append(N_r, new_r,1)

    N_fT = np.transpose(N_f)
    N_rT = np.transpose(N_r)
    N = N_r - N_f
    N_T = N_rT - N_fT

    num_cols = len(N[0])
    I = np.identity(num_cols)

    M = np.append(np.append(I, N_fT,1), np.append(I, N_rT,1),0)
    M_rref = sympy.Matrix(M).rref()

    # Set up the vectors for kinetic rate constants
    #  ['cAMP','LRGbinding_B1AR','LRGbinding_M2','B1AR','GsProtein','GiProtein']

    kf = []
    kr = []
    if combined_kinetic_parameters:
        [k_kinetic, N_cT, K_C, W] = grand_kinetic_parameters(M, True, dims, V, 0)
        sys_struct[sys_name]['kfkr'] = k_kinetic
        sys_struct[sys_name]['Kc'] = K_C
        sys_struct[sys_name]['N_cT'] = N_cT
        kf = k_kinetic[:num_cols]
        kr = k_kinetic[num_cols:]
    else:
        for sys_name in array_names:
            nrx = int(len(sys_struct[sys_name]['kfkr'])/2)
            kf.extend(sys_struct[sys_name]['kfkr'][:nrx])
            kr.extend(sys_struct[sys_name]['kfkr'][nrx:])

            k_kinetic = kf + kr

    W = list(np.append([1]*len(N[0]), [V_myo]*num_rows))

    lambda_expo = np.matmul(np.linalg.pinv(M), [math.log(k) for k in k_kinetic])
    lambdaW = [math.exp(l) for l in lambda_expo]
    lambdak = [lambdaW[i]/W[i] for i in range(len(W))]
    kappa = lambdak[:len(N[0])]
    K = lambdak[len(N[0]):]

    # save([output_dir 'FCU_adenylylCyclase_BG_parameters.mat'],'kappa','K','k_kinetic')

    # Checks
    N_rref = sympy.Matrix(N).rref()
    # R_mat = sympy.Matrix(N).nullspace() #'r'
    R_mat = null_space(N) #'r'

    k_est = np.matmul(M,[math.log(k) for k in lambdaW])
    k_est = [math.exp(k) for k in k_est]
    diff = [(k_kinetic[i] - k_est[i])/k_kinetic[i] for i in range(len(k_kinetic))]
    error = np.sum([abs(d) for d in diff])

    K_eq = [kf[i]/kr[i] for i in range(len(kr))]
    zero_est = np.matmul(np.transpose(R_mat),K_eq)
    zero_est_log = np.matmul(np.transpose(R_mat),[math.log(k) for k in K_eq])


    # ### print outputs ###
    for ik in range(len(kappa)):
        print('var kappa_%s: fmol_per_sec {init: %g, pub: out}' %(rxnIDs[ik],kappa[ik]))
    for ik in range(len(Kunique)):
        print('var K_%s: per_fmol {init: %g, pub: out}' %(Kunique[ik],K[ik]))

    print('error =', (error))

    # initialise struct for storing modules contributing to a given K
    for ik in range(len(Kunique)):
        Kname_modules[Kunique[ik]] = []

    for sys_name in array_names:
        modKname = sys_struct[sys_name]['Kname']
        for ik in range(len(modKname)):
            Kname_modules[modKname[ik]].append(sys_name)

    # write out CellML code
    if True:
        j = 10
        cellmlfilepath = os.getcwd() + '\\TEMP.cellml.txt'
        with open(cellmlfilepath,'w') as cid:
       
            cid.write('def model BG_muscarinic_composite as\n def import using "units_and_constants/units_BG.cellml" for\n\
            unit mM using unit mM;\nunit fmol using unit fmol;\nunit per_fmol using unit per_fmol;\n\
            unit J_per_mol using unit J_per_mol;\nunit fmol_per_sec using unit fmol_per_sec;\n\
            unit C_per_mol using unit C_per_mol;\n  unit J_per_C using unit J_per_C;\n\
            unit microm3 using unit microm3;\n  unit fF using unit fF;\n\
            unit fC using unit fC;\n  unit fA using unit fA;\n\
            unit per_second using unit per_second;\n  unit millivolt using unit millivolt;\n\
            unit per_sec using unit per_sec;\n  unit J_per_K_per_mol using unit J_per_K_per_mol;\n\
            unit fmol_per_L using unit fmol_per_L;\n  unit fmol_per_L_per_sec using unit fmol_per_L_per_sec;\n\
            unit per_sec_per_fmol_per_L using unit per_sec_per_fmol_per_L;\n  unit uM using unit uM;\n\
            unit mM_per_sec using unit mM_per_sec;\n  unit uM_per_sec using unit uM_per_sec;\n\
            unit pL using unit pL;\n  unit m_to_u using unit m_to_u;\n enddef;\n')
            cid.write('def import using "units_and_constants/constants_BG.cellml" for\n\
                comp constants using comp constants;\nenddef;\n\n')
            for module in array_names:
                cid.write('def import using "%s/BG_%s.cellml" for\ncomp %s using comp %s;\nenddef;\n' %(module,module,module,module))
            cid.write('\ndef comp BG_parameters as\n')
            for ik in range(len(kappa)):
                cid.write('var kappa_%s: fmol_per_sec {init: %g, pub: out};\n' %(rxnIDs[ik],kappa[ik]))
            for ik in range(len(Kunique)):
                cid.write('var K_%s: per_fmol {init: %g, pub: out};\n' %(Kunique[ik],K[ik]))
            cid.write('enddef;\n')
            cid.write('    def comp environment as\n\
                var time: second {pub: out};\n\
                var vol_myo: pL {init: 34.4, pub: out};\n\
                var freq: dimensionless {init: 500};\n\
                // stimulus\n\
                // ramp UP and ramp DOWN\n\
                var stimSt: second {init: 3.5e-4};\n\
                var stimDur: second {init: 0.25e-4};\n\
                var tR: second {init: 1.8e-4};\n\
                var stimMag: fmol {init: 1e1};\n\
                var stimHolding: fmol {init: 1e-5};  \n\
                var m: fmol_per_sec;  \n\
                m = stimMag/tR;   \n\
                q_L_M2_init = sel          \n\
                    case (time < stimSt) and (time > stimSt-tR): \n\
                        stimHolding+m*(time-stimSt+tR);   \n\
                    case (time >= stimSt) and (time < stimSt+stimDur):   \n\
                        stimMag+stimHolding;   \n\
                    case (time < stimSt+tR+stimDur) and (time >= stimSt+stimDur):\n\
                        stimHolding+-m*(time-stimSt-tR-stimDur);  \n\
                    otherwise:             \n\
                        stimHolding;       \n\
                endsel;\n')
            for j in range(len(K)):
                cid.write('var q_%s_init: fmol {init: 1e-888};\n' %Kunique[j])
            cid.write('\n// Global value\n')
            for j in range(len(K)):
                cid.write('var q_%s: fmol {pub: out};\n' %Kunique[j])
            for module in array_names:
                modKname = sys_struct[module]['Kname']
                cid.write('\n// %s imports\n' %module)
                for j in modKname:
                    cid.write('var q_%s_m%s: fmol {pub: in};\n' %(j, module))
                cid.write('\n')
            cid.write('\n')
            for kun in Kunique:
                cid.write('q_%s = q_%s_init'%(kun,kun))
                for mod in Kname_modules[kun]:
                    cid.write(' + q_%s_m%s ' %(kun,mod))
                cid.write(';\n')
            cid.write('enddef;\n')

            if False: # the below is for an individual module
                print('\n')
                print('// Input from global environment')
                for j in range(len(K)):
                    print('var q_%s_global: fmol {pub: in};\n' %Kunique[j])
                print('// Output to global environment')
                for j in range(len(K)):
                    print('var q_%s: fmol {init: 1e-16, pub: out};\n' %Kunique[j])
                for j in range(len(K)):
                    print('mu_%s = R*T*ln(K_%s*q_%s_global);\n' %(Kunique[j],Kunique[j],Kunique[j]))
            cid.write('\n')
            for module in array_names:
                modKname = sys_struct[module]['Kname']
                cid.write('def map between environment and %s for\n' %module)
                cid.write('vars time and time;\n')
                for mod in modKname:
                    cid.write('vars q_%s_m%s and q_%s;\n' %(mod, module,mod))
                    cid.write('vars q_%s and q_%s_global;\n' %(mod, mod))
                cid.write('enddef;\n\n')

            for module in array_names:
                modKname = sys_struct[module]['Kname']
                modrxnID = sys_struct[module]['rxnID']
                cid.write('def map between BG_parameters and %s for\n' %(module))
                for ik in modrxnID:
                    cid.write('vars kappa_%s and kappa_%s;\n'%(ik,ik))
                for mod in modKname:
                    cid.write('vars K_%s and K_%s;\n' %(mod, mod))
                cid.write('enddef;\n')
            cid.write('\n')
            for module in array_names:
                cid.write('def map between constants and %s for\n' %(module))
                cid.write('\tvars R and R;\n\tvars T and T;\nenddef;\n')

            cid.write('\nenddef;\n')
        cid.close()
    #
    #