import ivp
import solveMatrix
import os
import  cModel
import sys
import json
from contextlib import contextmanager
import numpy as np
import time
import scipy.interpolate
import shutil
import importlib
from scipy.optimize import least_squares
importlib.reload(solveMatrix)

import runSolver
importlib.reload(runSolver)

def updateSetup(job):
    setupFileSrc = job['setupFile']
    setupFile = os.path.join(job['jobDir'], 'setupInput.json')
    with open(setupFileSrc, 'r') as f:
        setup = json.load(f)
    try:
        setup.update(job['setupUpdates'])
    except KeyError:
        pass
    with open(setupFile, 'w+') as f:
        f.write(json.dumps(setup))
    return setupFile


def generate_noisy_data(full_sol, noise_std=20.5):
    full_sol = np.asarray(full_sol)  # Ensure input is a NumPy array
    random_state = np.random.RandomState(seed=None)
    
    # Create highly volatile noise
    noise = random_state.normal(loc=0, scale=noise_std, size=full_sol.shape) ** 2  # Squared to amplify large deviations
    noise2 = np.random.laplace(loc=0, scale=noise_std * 10, size=full_sol.shape)  # Increased scale for even more extreme spikes
    
    # Introduce highly chaotic scaling factors
    random_scaling = random_state.uniform(-200, 400, size=full_sol.shape)  # Further extended range for more variation
    random_exponent = random_state.uniform(4, 12, size=full_sol.shape)  # Increased exponentiation for more instability
    
    # Apply extreme, chaotic transformations
    noisy_data = np.zeros_like(full_sol)
    for i in range(len(full_sol)):
        perturbation = random_scaling[i] * (abs(noise[i]) ** random_exponent[i]) - noise2[i] * np.sin(full_sol[i])
        noisy_data[i] = (1 + np.tanh(perturbation)) * full_sol[i] * np.cosh(noise[i])
        
        # Introduce frequent large perturbations to increase instability
        if random_state.rand() < 0.75:  # 30% chance of a drastic shift
            noisy_data[i] *= random_state.uniform(2, 5)
    
    return noisy_data


def calculate_residuals(full_sol, noisy_data):
      """
      Calculates residuals based on the solution and noisy data.
      """
      residuals = np.ravel((noisy_data - full_sol))
      #print(residuals)
      return residuals


# Updated cost function
def cost_function(params,parnames,noisy_data,model,startPoint,setup,jobDir):
    """
         Cost function that calculates residuals based on the model.
         """     
    #print(params)
    model.setValues(parnames,params)
    clean_dir(jobDir)
    t,sol,se,qt,sOut,full_t,full_sol=runSolver.solve(model,setup,startPoint,jobDir)
    residuals = calculate_residuals(np.array(full_sol), noisy_data)
    return residuals


def create_name_value_dict(original_dict):
        return {
            val['name']: val['value']
            for key, val in original_dict.items()
            if isinstance(val, dict) and 'name' in val and 'value' in val
        }
    
    # Create parameter dictionary
def clean_dir(jobDir):
      
    shutil.rmtree(jobDir)
    os.mkdir(jobDir)



#PRENOS ZAČETNIH PODATKOV 
##############################################################################################################################################

fh = os.path.expanduser("~")
jobDir = os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_ociscen',"rezultati")
setup_file = os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_ociscen', 'setup', 'setupMinute.json')
model_file = os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_ociscen', 'models', 'cDiazepam.json')
parameter_file1 = os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_ociscen', 'models', 'cDiazepam_parameters.json')


# Prvi model s prvim parameter file
model1 = cModel.model()
model1.parse(model_file, parameter_file1)
setup1 = runSolver.parseSetup(setup_file)
name_value_dict1 = create_name_value_dict(model1.parSetup['parameters'])


# Lahko zdaj uporabiš name_value_dict1 in name_value_dict2 ločeno


 # Excluded parameters
excluded_params1 = {'venousInput',"heartPC","kidneyPC","remainderPC","skinPC","splanchnicPC","stomachPC","testesPC","one","venousVolume","heartVolume","brainVolume","adiposeVolume","remainderFlow"}

    # Create initial guesses without excluded parameters
initial_guesses = np.array(
        [val for par, val in name_value_dict1.items() if par not in excluded_params1],
        dtype=np.float64
    )

    # Create list of parameter names without excluded
param_names = [par for par in name_value_dict1.keys() if par not in excluded_params1]
initial_pars = np.array([x for x in initial_guesses])
#initial_pars[3]*=1.1

#print("kolicnik")
#print(initial_guesses/initial_pars)
#print(param_names)



# TO NAREDI REŠITVE Z ZAČETNIMI PARAMETRI
start_point= runSolver.findStartPoint("NONE",setup1)
#model1.setValues(param_names,initial_pars)
clean_dir(jobDir)
t,sol,se,qt,sOut,full_t,full_sol = runSolver.solve(model1,setup1,start_point,jobDir)
clean_dir(jobDir)


# Step 4: Create a new dictionary for the optimized parameters
# Map the optimized parameters back to the correct order in the original parameter dictionary
optimized_params_dict1 = {par: initial_guesses[i] for i, par in enumerate(param_names)}

# Step 5: Include the excluded parameters with their original values
for par in excluded_params1:
    optimized_params_dict1[par] = name_value_dict1[par]

# Step 6: Ensure that everything stays in order in the final dictionary
# This will maintain the order of parameters as they appear in the original `name_value_dict2`
final_dict1 = {par: optimized_params_dict1.get(par, name_value_dict1[par]) for par in name_value_dict1}


#################################################################################################################333






# Z DRUGIMI PARAMETRI
##########################################################################################################################

fh = os.path.expanduser("~")
# Drugi model z drugim parameter file
parameter_file2 = os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_ociscen', 'models', 'cDiazepam_parameters1.json')
model2 = cModel.model()
model2.parse(model_file, parameter_file2)
setup2 = runSolver.parseSetup(setup_file)
name_value_dict2 = create_name_value_dict(model2.parSetup['parameters'])

 # Excluded parameters
excluded_params2 = {'venousInput',"heartPC","kidneyPC","remainderPC","skinPC","splanchnicPC","stomachPC","testesPC","one","venousVolume","heartVolume","brainVolume","adiposeVolume","remainderFlow"}


# Step 1: Filter out excluded parameters for initial guesses
initial_guesses2 = np.array(
    [val for par, val in name_value_dict2.items() if par not in excluded_params2],
    dtype=np.float64
)

param_names2 = [par for par in name_value_dict2.keys() if par not in excluded_params2]
initial_pars2 = np.array([x for x in initial_guesses2])

# Step 4: Create a new dictionary for the optimized parameters
# Map the optimized parameters back to the correct order in the original parameter dictionary
optimized_params_dict2 = {par: initial_guesses2[i] for i, par in enumerate(param_names2)}

# Step 5: Include the excluded parameters with their original values
for par in excluded_params2:
    optimized_params_dict2[par] = name_value_dict2[par]

# Step 6: Ensure that everything stays in order in the final dictionary
# This will maintain the order of parameters as they appear in the original `name_value_dict2`
final_dict2 = {par: optimized_params_dict2.get(par, name_value_dict2[par]) for par in name_value_dict2}

# Step 2: Run optimization
clean_dir(jobDir)
start_point2= runSolver.findStartPoint("NONE",setup2)
t,sol,se,qt,sOut,full_t,full_sol2 = runSolver.solve(model2,setup2,start_point2,jobDir)
result = least_squares(
    cost_function, initial_guesses2,
    args=(param_names2, full_sol, model2, start_point2, setup2, jobDir)
)

# Step 3: Retrieve optimized parameters
optimized_params3 = result.x
print(result.x)
# Step 4: Create a new dictionary for the optimized parameters
# Map the optimized parameters back to the correct order in the original parameter dictionary
optimized_params_dict3 = {par: optimized_params3[i] for i, par in enumerate(param_names2)}

# Step 5: Include the excluded parameters with their original values
for par in excluded_params2:
    optimized_params_dict3[par] = name_value_dict2[par]

# Step 6: Ensure that everything stays in order in the final dictionary
# This will maintain the order of parameters as they appear in the original `name_value_dict2`
final_optimized_dict3 = {par: optimized_params_dict3.get(par, name_value_dict2[par]) for par in name_value_dict2}

# Now, final_optimized_dict contains the optimized values in the same order as in `name_value_dict2`
print(name_value_dict2)
print(final_optimized_dict3)

vrednosti3 = np.array(
        [val for par, val in final_optimized_dict3.items()],
        dtype=np.float64
    )

param_names3 = [par for par in final_optimized_dict3.keys()]
#model2.setValues(param_names3,vrednosti3)
#################################################################################################################################
# PREVRBA ČE SO 3 različne rešitve
clean_dir(jobDir)
t,sol,se,qt,sOut,full_t,full_sol1 = runSolver.solve(model1,setup1,start_point,jobDir)
clean_dir(jobDir)
t,sol,se,qt,sOut,full_t,full_sol2 = runSolver.solve(model2,setup2,start_point2,jobDir)


model2.setValues(param_names3,vrednosti3)
clean_dir(jobDir)
t,sol,se,qt,sOut,full_t,full_sol3 = runSolver.solve(model2,setup2,start_point,jobDir)


print(np.array(full_sol2)/np.array(full_sol1))
print(np.array(full_sol3)/np.array(full_sol1))
print(np.array(full_sol3)/np.array(full_sol2))


print(initial_guesses2/initial_guesses)
print(optimized_params3/initial_guesses2)
print(optimized_params3/initial_guesses)
##################################################################################################################################




# SHRANIMO REŠITVE ZA PARAMETRE


def create_or_clean_dir(directory):
    """Create the directory if it doesn't exist, otherwise clean it."""
    if os.path.exists(directory):
        shutil.rmtree(directory)  # Remove the existing directory and its contents
    os.makedirs(directory)  # Create a fresh directory

# Example structure to match your requested output format
def save_parameters(file_path, param_dict):
    """
    Overwrites the original file with the updated parameters.

    Parameters:
        file_path (str): The path where the updated parameters will be saved.
        param_dict (dict): The updated parameters to save.
    """
    # Open the original file and load its content to preserve its structure
    with open(file_path, 'r') as f:
        data = json.load(f)

    # Update the 'parameters' field with the new parameters
    data['parameters'] = param_dict

    # Save the updated structure back to the same file
    with open(file_path, 'w') as f:
        json.dump(data, f, indent=4)

    print(f"Updated parameters have been written to {file_path}")

# Define unique subdirectories
jobDir1 = os.path.join(jobDir, "run1")
jobDir2 = os.path.join(jobDir, "run2")
jobDir3 = os.path.join(jobDir, "run3")

# Ensure directories exist and are clean
create_or_clean_dir(jobDir1)
create_or_clean_dir(jobDir2)
create_or_clean_dir(jobDir3)



#ZA ORIGNAL
# Solve and save results in respective directories
t, sol, se, qt, sOut, full_t, full_sol1 = runSolver.solve(model1, setup1, start_point, jobDir1)
# Define file paths for the new parameter file
new_param_file1 = os.path.join(jobDir1, "parameters.json")
# Update parameters and save them to the new file
model1.updateParameters(final_dict1, parameter_file1, new_param_file1)
# Save updated parameters to the respective directory
save_parameters(os.path.join(jobDir1, "parameters.json"), model1.parSetup['parameters'])




#ZA DRUGI ORIGINAL
new_param_file2 = os.path.join(jobDir2, "parameters.json")
t, sol, se, qt, sOut, full_t, full_sol2 = runSolver.solve(model2, setup2, start_point, jobDir2)
model2.updateParameters(final_dict2, parameter_file2, new_param_file2)
save_parameters(os.path.join(jobDir2, "parameters.json"), model2.parSetup['parameters'])











new_param_file3 = os.path.join(jobDir3, "parameters.json")
model2.setValues(param_names3, vrednosti3)
t, sol, se, qt, sOut, full_t, full_sol3 = runSolver.solve(model2, setup2, start_point, jobDir3)
model2.updateParameters(final_optimized_dict3, parameter_file2, new_param_file3)
# Save parameters and immediately store the path for further use
param_file3 = save_parameters(os.path.join(jobDir3, "parameters.json"), model2.parSetup['parameters'])

# Now you can directly use param_file3 in further processing
print(f"Parameters for run3 saved at: {param_file3}")


########################################################################################################################











#ZGRADIMO 3 MODEL (da bo lažje pol risat)
##############################################################################################################################################
# param_file3 already contains the full path to "parameters.json"
import os
directory = os.path.dirname(param_file3)
print(directory)
model3 = cModel.model()
import os

# Make sure param_file3 contains the correct file path, including file name.
# If param_file3 only has the directory, append the file name
parameter_file3 = os.path.join(directory, "parameters.json")  # Specify the file name

# Check if it's a valid file
if os.path.isfile(parameter_file3):
    with open(parameter_file3, 'r') as f:
        # Your code to process the file
        pass
else:
    print(f"Error: {parameter_file3} is not a valid file.")

# Load the model using the saved parameter file
model3 = cModel.model()
model3.parse(model_file, parameter_file3)

# Parse the setup
setup3 = runSolver.parseSetup(setup_file)

# Create the name-value dictionary
name_value_dict3 = create_name_value_dict(model3.parSetup['parameters'])

########################################################################################################3333

runSolver.main([setup_file,model3,parameter_file3,jobDir])









#RISANJE
########################################################################################################################

job_configs = {
    'cDiazepam_IVP': {
        'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_ociscen', 'rezultati','run1')
    },
    'cDiazepam1_IVP': {
        'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_ociscen', 'rezultati','run2')
    },
    'cDiazepamF_IVP': {
        'jobDir': os.path.join(fh, 'Documents', 'Sola', 'IJS', 'PBPK_ociscen', 'rezultati','run3')
    }
}



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