import numpy as np
import matplotlib.pyplot as plt
import os
import json
import cModel
import runSolver
import importlib

importlib.reload(cModel)
importlib.reload(runSolver)

# Load data
def load_data(modelName, jobDir):
    Q = runSolver.loadSolutionFromDir(jobDir, True)
    return Q

def print_lut_keys(Q):
    lut = Q['lut']
    # Debug information - removed print statements
    # print("Available keys in lut:")
    # for key in lut:
    #     print(f"Key: {key}, Index: {lut[key]}")

def plot_rbc_to_plasma_ratio(Q):
    lut = Q['lut']
    
    # Debug information - removed print statements
    # print("Available keys in lut:", lut.keys())
    
    # Example keys based on available compartments
    red_blood_cells_key = 'arterial'
    plasma_key = 'venous'
    
    if red_blood_cells_key not in lut or plasma_key not in lut:
        raise KeyError(f"Expected keys '{red_blood_cells_key}' or '{plasma_key}' are missing in 'lut'.")
    
    fy1 = Q['sol'][:, lut[red_blood_cells_key]]
    fy = Q['sol'][:, lut[plasma_key]]
    qy = fy1[1:] / fy[1:]
    t = Q['t']
    
    plt.plot(t[1:], qy)
    plt.xlabel('Time')
    plt.ylabel('RBC/Plasma Ratio')
    plt.title('RBC to Plasma Ratio Over Time')
    plt.show()

# Parse model and retrieve parameters
def parse_and_retrieve_parameters(modelFile, parameterFile):
    model = cModel.model()
    model.parse(modelFile, parameterFile)
    # Debug information - removed print statement
    # print('w(oI) {}'.format(model.getWeight('oralIngestion')))
    return model

def calculate_derivatives_weights(model, sOut, lut, compartment):
    """
    Calculate derivatives and weights for a specific compartment and print sorted results.

    Parameters:
    - model: An instance of the model object
    - sOut: Simulation output data
    - lut: Look-up table (LUT) of compartments
    - compartment: The compartment for which to calculate derivatives and weights
    """
    try:
        # Compute derivatives for the specified compartment
        d = model.getDerivatives(sOut, lut[compartment])
        print(sOut)
        print(lut[compartment])
        # Retrieve the LUT for parameter indices and weights
        lutSE = model.lutSE  # Accessing lutSE as an attribute
        w = model.getWeights(lutSE)
        print(lutSE)
        print(w)
        print(d)
        # Ensure the lengths of arrays match
        if len(lutSE) != len(d) or len(lutSE) != len(w):
            raise ValueError("Mismatch in lengths of LUT, derivatives, or weights.")
        
        # Calculate weighted derivatives
        s = {x: d[lutSE[x]] * w[lutSE[x]] for x in lutSE}
        s = dict(sorted(s.items(), key=lambda item: item[1], reverse=True))
        # Print the results
        print(f"Top 10 derivatives and weights for compartment '{compartment}':")
        for p in list(s)[:10]:
            j = lutSE[p]
            print(f'{p}: Weighted Derivative = {d[j] * w[j]:.2g} (Derivative = {d[j]:.2g}, Weight = {w[j]:.2g})')
    
    except AttributeError as e:
        print(f"Error accessing model methods or attributes: {e}")
    
    except ValueError as e:
        print(f"Value error: {e}")



def retrieve_and_calculate_total_mass(model, setup):
    tscale = runSolver.getScale(setup)
    
    # Attempt to retrieve parameters with attribute checking
    blood_to_plasma_pc_scale = getattr(model, 'bloodToPlasmaPCscale', None)
    blood_volume = getattr(model, 'bloodVolume', None)
    
    # Print debug information
    if blood_to_plasma_pc_scale is not None:
        print(f'Blood to Plasma PC Scale: {blood_to_plasma_pc_scale}')
    else:
        print('Blood to Plasma PC Scale parameter is missing.')
    
    if blood_volume is not None:
        print(f'Blood Volume: {blood_volume}')
    else:
        print('Blood Volume parameter is missing.')
    
    return blood_to_plasma_pc_scale, blood_volume




def main():     #PAZI MODEL in PArameter File morta bit ista!!!
    # Define job directory and model files
    jobDir = os.path.join(os.path.expanduser('~'), 'Documents', 'Sola', 'IJS', 'PBPK_public', 'cDiazepam1_IVP')
    modelFile = os.path.join(os.path.expanduser('~'), 'Documents', 'Sola', 'IJS', 'PBPK_public', 'models', 'cDiazepam.json')
    parameterFile = os.path.join(os.path.expanduser('~'), 'Documents', 'Sola', 'IJS', 'PBPK_public', 'models', 'cDiazepam_parameters.json')
    
    # Load data
    Q = load_data('cDiazepam1_IVP', jobDir)
    
    # Plot RBC to Plasma Ratio
    plot_rbc_to_plasma_ratio(Q)
    
    # Parse model and retrieve parameters
    model = parse_and_retrieve_parameters(modelFile, parameterFile)
    
    # Calculate derivatives and weights for 'kidney'
    compartment = 'kidney'
    calculate_derivatives_weights(model, Q['sOut'], Q['lut'], compartment)
    
    # Retrieve model and calculate total mass
    setup = Q['setup']
    retrieve_and_calculate_total_mass(model, setup)

if __name__ == "__main__":
    main()