import numpy as np
import matplotlib.pyplot as plt
import os
import cModel
import time
import pdfkit
import tempfile
import runSolver
import io
import importlib
#Nariše graf iz IVP modelov
importlib.reload(cModel)
importlib.reload(runSolver)

def getModel(solution, modelName):
    """
    Load and parse the model from the given solution dictionary using the modelName.
    """
    Q = solution[modelName]  # Retrieve the solution dictionary for the specified model
    model = cModel.model()  # Create an instance of the cModel class
    setupFile = Q['setup']  # Get the setup file path from the solution
    modelFile = Q['model']  # Get the model file path from the solution
    parameterFile = Q['parameters']  # Get the parameter file path from the solution
    print('modelFile: {} {}'.format(modelFile, os.path.isfile(modelFile)))  # Print the model file path and its existence status
    print('parameterFile: {} {}'.format(parameterFile, os.path.isfile(parameterFile)))  # Print the parameter file path and its existence status
    model.parse(modelFile, parameterFile)  # Parse the model file with the given parameters
    return model  # Return the loaded model

def mergeSolutions(seq):
    """
    Merge multiple solution dictionaries into a single dictionary.
    """
    out = {}
    for v in ['t', 'sol', 'se', 'qt', 'sOut']:  # Merge arrays from each solution
        out[v] = np.concatenate([x[v] for x in seq])
    for v in ['lut', 'lutSE', 'setup', 'model', 'parameters', 'qt', 'sOut']:  # Use data from the last solution for these keys
        out[v] = seq[-1][v]
    return out  # Return the merged dictionary

# Define paths
fh = os.path.expanduser('~')

# Define job configurations for loading
job_configs = {
    'cDiazepam_IVP': {
        'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_public', 'cDiazepam_IVP')
    },
    'cDiazepam1_IVP': {
        'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_public', 'cDiazepam1_IVP')
    },
    'cDiazepamF_IVP': {
        'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_public', 'cDiazepamF_IVP')
    },
    'cDiazepamB_IVP': {
        'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_public', 'cDiazepamB_IVP')
    }
}

# Load and merge solutions
solution = {}
for modelName, job in job_configs.items():
    seq = [runSolver.loadSolutionFromDir(job['jobDir'], True)]
    solution[modelName] = mergeSolutions(seq)

def analyze_model(model, setup):
    """
    Perform matrix operations and print results.
    """
    tscale = runSolver.getScale(setup)  # Get the time scale from the setup
    print(f'tscale={tscale}')  # Print the time scale
    model.inspect()  # Inspect the model
    
    print("***********done************")
    print(model.M(1).shape)  # Print the shape of the matrix from the model
    print(model.m)  # Print the model parameters

    nt = setup['nt']  # Number of time points
    qtmax = 1  # Maximum time (minute)
    qt = np.linspace(0, qtmax, nt)  # Generate time points
    par = 'venousInput'  # Parameter to plot

    # Plot parameters
    try:
        hw = model.get(par)  # Get the parameter from the model
        print(hw)  # Print the parameter information
        ht = [10 * hw['value'](x) for x in qt]  # Calculate parameter values over time
        plt.plot(qt / tscale, ht)  # Plot the parameter values
    except (KeyError, TypeError):  # Handle errors if the parameter is not found
        print(f'Troubles getting {par}')
        pass

    # Measure performance of matrix operations
    start_time = time.time()  # Record start time
    for i in range(100000):  # Perform matrix operations 100,000 times
        model.M(1e7)  # Call the matrix operation
    end_time = time.time()  # Record end time
    print('Time: {:.3f} s'.format(end_time - start_time))  # Print elapsed time

    fM = model.M(0)  # Get the matrix from the model
    print('Rank: {}/{}'.format(np.linalg.matrix_rank(fM), fM.shape))  # Print the rank of the matrix

    np.set_printoptions(suppress=True, precision=2, linewidth=150)  # Set NumPy print options
    print(f'{fM}')  # Print the matrix

    v, Q = np.linalg.eig(fM)  # Perform eigenvalue decomposition
    np.set_printoptions(suppress=False)  # Reset NumPy print options
    print(Q[:, 2:4])  # Print selected eigenvectors

    Q1 = np.linalg.inv(Q)  # Calculate the inverse of the eigenvector matrix
    D = np.diag(v)  # Create a diagonal matrix of eigenvalues

    t = np.linspace(0, 100, 101)  # Generate time points for plotting
    fy = [model.u(x)[14] for x in t]  # Calculate model output over time
    print('{} {}'.format(len(fy), len(t)))  # Print lengths of output and time points
    plt.figure()  # Create a new figure
    plt.plot(fy)  # Plot the model output
    fu = model.u(0)  # Get the model output at time 0
    print(Q1 @ fu)  # Print the result of the matrix multiplication

def plot_results(solution, modelName):
    """
    Plot results for the specified model from the solution.
    """
    model = getModel(solution, modelName)  # Load the model
    setup = solution[modelName]['setup']  # Get the setup from the solution
    tscale = runSolver.getScale(setup)  # Get the time scale
    print(f'tscale={tscale}')  # Print the time scale
    model.inspect()  # Inspect the model

    name = ['venous', 'adipose', 'brain', 'heart', 'kidney', 'liver', 'lung', 'muscle', 'skin', 'stomach', 'splanchnic', 'excrement']  # Compartment names
    tmax = solution[modelName]['t'][-1]  # Get the maximum time from the solution

    # Define colors and shading for models
    models = {
        'cDiazepam_IVP': {'color': 'orange', 'shadeColor': 'red'},
        'cDiazepamF_IVP': {'color': 'blue', 'shadeColor': 'green'},
        'cDiazepam1_IVP': {'color': 'black', 'shadeColor': 'yellow'},
        'cDiazepamB_IVP': {'color': 'purple', 'shadeColor': 'gold'}
    }

    fig, axs = plt.subplots(4, 3, figsize=(15, 20))  # Create a grid of subplots

    for i in range(len(name)):  # Loop over each compartment
        row = i // 3  # Determine row index
        col = i % 3  # Determine column index
        ax = axs[row, col]  # Get the subplot

        for m in models:  # Loop over each model
            fM = models[m]  # Get model configuration
            Q = solution[m]  # Get solution for the model
            v = name[i]  # Get the compartment name

            try:
                j = Q['lut'][v]  # Get the index for the compartment
            except KeyError:
                try:
                    v1 = alias[v]  # Handle alias if needed
                    j = Q['lut'][v1]  # Get the index for the alias
                except KeyError:
                    print(f'No data for {v}/{v1}')  # Print error if compartment not found
                    continue

            fy = Q['sol'][:, j]  # Get the solution for the compartment
            fe = Q['se'][:, j]  # Get the standard error for the compartment
            t = Q['t']  # Get the time points

            # Plot the mean values and confidence intervals
            ax.plot(t / tscale, fy, color=fM['color'], label=f'{m} Mean')
            ax.fill_between(t / tscale, fy - fe, fy + fe, color=fM['shadeColor'], alpha=0.1)  # Shaded area for confidence intervals
            ax.plot(t / tscale, fy - fe, color=fM['shadeColor'], linewidth=1, alpha=0.2)  # Lower bound
            ax.plot(t / tscale, fy + fe, color=fM['shadeColor'], linewidth=1, alpha=0.2)  # Upper bound

        # Find the maximum y-value including confidence intervals
        max_y_value = max(np.max(fy + fe), np.max(fy - fe))  
        ax.set_ylim([0, min(max_y_value, 10000)])  # Set y-axis limit to max value or 5500, whichever is smaller
        ax.set_xlim([0, 1.1 * tmax / tscale])  # Set x-axis limits
        ax.set_xlabel(setup['tUnit'])  # Set x-axis label
        ax.set_title(name[i])  # Set plot title
        ax.legend()  # Add legend

    output_file = os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_public', 'results', 'plot.png')  # Define output file path
    plt.savefig(output_file)  # Save the plot to file
    plt.show()  # Display the plot


def main():
    """
    Main function to load data, analyze model, and plot results.
    """
    fh = os.path.expanduser('~')  # Define file path root

    # Define job configurations for loading
    job_configs = {
        'cDiazepam_IVP': {
            'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_public', 'cDiazepam_IVP')
        },
        'cDiazepam1_IVP': {
            'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_public', 'cDiazepam1_IVP')
        },
        'cDiazepamF_IVP': {
            'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_public', 'cDiazepamF_IVP')
        },
        'cDiazepamB_IVP': {
            'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_public', 'cDiazepamB_IVP')
        }
    }

    # Load and merge solutions
    solution = {}
    for modelName, job in job_configs.items():
        seq = [runSolver.loadSolutionFromDir(job['jobDir'], True)]  # Load solution from directory
        solution[modelName] = mergeSolutions(seq)  # Merge solutions
    # Capture print output
    output_stream = io.StringIO()
    analyze_model(getModel(solution, 'cDiazepamF_IVP'), solution['cDiazepamF_IVP']['setup'])  # Analyze the model
    analysis_results = analyze_model(getModel(solution, 'cDiazepamF_IVP'), solution['cDiazepamF_IVP']['setup'])
    plot_results(solution, 'cDiazepamF_IVP')  # Plot the results

    

if __name__ == "__main__":
    main()  # Execute the main function