# This model builder will refer to `scripts/DSS_history_match/106r2_SingleFracScanner_test_func.py`
# Which provides a better, more stable baseline model for history matching tasks.
# This is working correctly now. I'm happy with it. :P
# Shenyao Jin, 11/24/2025

import numpy as np
import os
import datetime
import matplotlib.pyplot as plt
from fiberis.analyzer.Data1D.Data1D_Gauge import Data1DGauge
from fiberis.analyzer.Data2D.core2D import Data2D
from typing import List
from fiberis.moose.config import (
    MatrixConfig, SRVConfig, HydraulicFractureConfig, ZoneMaterialProperties, SimpleFluidPropertiesConfig,
    PointValueSamplerConfig, LineValueSamplerConfig, TimeSequenceStepper, InitialConditionConfig
)
from fiberis.moose.model_builder import ModelBuilder
from fiberis.io.reader_moose_vpp import MOOSEVectorPostProcessorReader

def build_baseline_model(**kwargs) -> ModelBuilder:
    """
    In this function, I extract the common baseline model building steps from 106r2_SingleFracScanner_test_func.py
    to create a reusable baseline model generator. This script should be much more stable and no mistakes should be made
    by QCing the original script multiple times.
    --Shenyao Jin

    :param project_name: The name of the project. Defaults to "BaselineModel".
    :param model_width: The width of the model domain in meters. Defaults to 200.0 * 0.3048.
    :param fracture_y_coords: A list of y-coordinates for the fractures. Defaults to [0.0].
    :param model_length: The length of the model domain in meters. Defaults to 800.0 * 0.3048.
    :param nx: The number of elements in the x-direction. Defaults to 200.
    :param ny_per_layer_half: The number of elements in the y-direction for each half-layer. Defaults to 80.
    :param bias_y: The mesh bias in the y-direction. Defaults to 1.2.
    :param matrix_perm: The permeability of the matrix. Defaults to 1e-18.
    :param srv_perm: The permeability of the SRV. Defaults to 1e-15.
    :param fracture_perm: The permeability of the fracture. Defaults to 1e-13.
    :param srv_length_ft: The length of the SRV in feet. Defaults to 400.
    :param srv_height_ft: The height of the SRV in feet. Defaults to 20.
    :param hf_length_ft: The length of the hydraulic fracture in feet. Defaults to 250.
    :param hf_height_ft: The height of the hydraulic fracture in feet. Defaults to 0.2.
    :param initial_pressure: The initial pressure of the model. Defaults to 5.17E7.
    :param monitoring_point_shift_ft: The shift of the monitoring point in feet. Defaults to 80.
    :param start_offset_y: The start offset of the fiber sampler in the y-direction. Defaults to 20.
    :param end_offset_y: The end offset of the fiber sampler in the y-direction. Defaults to 20.
    :param num_fiber_points: The number of points in the fiber sampler. Defaults to 200.
    :return: A ModelBuilder object representing the baseline model.
    """
    # Define default parameters
    conversion_factor = 0.3048  # feet to meters

    # "data/fiberis_format/post_processing/injection_pressure_full_profile.npz" <- injection pressure profile
    # Load gauge data for MOOSE, I have already packed the data in fiberis format.
    gauge_data_for_moose = Data1DGauge()
    gauge_data_for_moose.load_npz("data/fiberis_format/prod/gauges/pressure_g1.npz")

    # Load DSS data so that I can crop the time range accordingly.
    DSS_data = Data2D()
    DSS_data.load_npz("data/fiberis_format/s_well/dss_data/Mariner 14x-36-POW-S - RFS strain change.npz")
    gauge_data_for_moose.select_time(DSS_data.start_time, DSS_data.get_end_time())
    gauge_data_for_moose.adaptive_downsample(130)
    gauge_data_for_moose.data = 6894.76 * gauge_data_for_moose.data  # Convert psi to Pa

    # Start building the model
    builder = ModelBuilder(project_name=kwargs.get("project_name", "BaselineModel"))
    domain_bounds = (- kwargs.get('model_width', 200.0 * conversion_factor),
                     + kwargs.get('model_width', 200.0 * conversion_factor))

    frac_coords = kwargs.get('fracture_y_coords', [0.0 * conversion_factor])
    # If frac_coords is a list (for multi-fracture mesh), use the first element for this single-fracture model's center.
    frac_y_center = frac_coords[0] if isinstance(frac_coords, list) else frac_coords
    domain_length = kwargs.get('model_length', 800.0 * conversion_factor)
    builder.build_stitched_mesh_for_fractures(
        fracture_y_coords=frac_coords,
        domain_bounds=domain_bounds,
        domain_length=domain_length,
        nx=kwargs.get('nx', 200),
        ny_per_layer_half=kwargs.get('ny_per_layer_half', 100),
        bias_y=kwargs.get('bias_y', 1.1)
    )

    matrix_perm = kwargs.get('matrix_perm', 1e-20)
    srv_perm = kwargs.get('srv_perm', 1e-16)
    fracture_perm = kwargs.get('fracture_perm', 1e-13)

    # The tensor format for permeability in fiberis:
    matrix_perm_str = f"{matrix_perm} 0 0 0 {matrix_perm} 0 0 0 {matrix_perm}"
    srv_perm_str = f"{srv_perm} 0 0 0 {srv_perm} 0 0 0 {srv_perm}"
    fracture_perm_str = f"{fracture_perm} 0 0 0 {fracture_perm} 0 0 0 {fracture_perm}"

    # Material properties
    matrix_mats = ZoneMaterialProperties(porosity=0.03, permeability=matrix_perm_str)
    srv_mats = ZoneMaterialProperties(porosity=0.032, permeability=srv_perm_str) # <- Changed here. Physically more reasonable.
    fracture_mats = ZoneMaterialProperties(porosity=0.16, permeability=fracture_perm_str)

    # Define Initial Conditions
    initial_pressure_val = kwargs.get('initial_pressure', 5.17E7)
    initial_pressure_val_srv = kwargs.get('initial_pressure_srv', gauge_data_for_moose.data[0])
    pressure_ic = InitialConditionConfig(
        name="initial_pressure",
        ic_type="ConstantIC",
        variable="pp",
        params={"value": initial_pressure_val}
    )

    pressure_ic_srv_frac = InitialConditionConfig(
        name="initial_pressure_srv_frac",
        ic_type="ConstantIC",
        variable="pp",
        params={"value": initial_pressure_val_srv}
    )

    builder.set_matrix_config(MatrixConfig(name="matrix", materials=matrix_mats, initial_conditions=[pressure_ic_srv_frac])) # changed here

    center_x_val = domain_length / 2.0
    srv_length_ft = kwargs.get('srv_length_ft', 400)
    srv_height_ft = kwargs.get('srv_height_ft', 5)  # <- Changed here. From 50 to 20
    hf_length_ft = kwargs.get('hf_length_ft', 250)
    hf_height_ft = kwargs.get('hf_height_ft', 0.2)

    geometries = [
        SRVConfig(name="srv", length=srv_length_ft * conversion_factor, height=srv_height_ft * conversion_factor,
                  center_x=center_x_val, center_y=frac_y_center, materials=srv_mats, initial_conditions=[pressure_ic_srv_frac]),
        HydraulicFractureConfig(name="hf", length=hf_length_ft * conversion_factor,
                                height=hf_height_ft * conversion_factor, center_x=center_x_val,
                                center_y=frac_y_center, materials=fracture_mats, initial_conditions=[pressure_ic_srv_frac])
    ]

    sorted_geometries = sorted(geometries, key=lambda x: x.height, reverse=True)
    next_block_id = 1
    for geom_config in sorted_geometries:
        if isinstance(geom_config, SRVConfig):
            builder.add_srv_config(geom_config)
            builder.add_srv_zone_2d(geom_config, target_block_id=next_block_id)
        elif isinstance(geom_config, HydraulicFractureConfig):
            builder.add_fracture_config(geom_config)
            builder.add_hydraulic_fracture_2d(geom_config, target_block_id=next_block_id)
        next_block_id += 1

    builder.add_nodeset_by_coord(nodeset_op_name="injection", new_boundary_name="injection",
                                 coordinates=(center_x_val, frac_y_center, 0))

    builder.add_variables([
        "pp",
        {"name": "disp_x", "params": {"initial_condition": 0}},
        {"name": "disp_y", "params": {"initial_condition": 0}}
    ])

    builder.set_porous_flow_dictator(dictator_name="dictator", porous_flow_variables="pp")
    builder.add_global_params({"PorousFlowDictator": "dictator", "displacements": "'disp_x disp_y'"})

    biot_coeff = 0.7
    builder.add_porous_flow_mass_time_derivative_kernel(kernel_name="dt", variable="pp")
    builder.add_porous_flow_darcy_base_kernel(kernel_name="flux", variable="pp")
    builder.add_stress_divergence_tensor_kernel(kernel_name="grad_stress_x", variable="disp_x", component=0)
    builder.add_stress_divergence_tensor_kernel(kernel_name="grad_stress_y", variable="disp_y", component=1)
    builder.add_porous_flow_effective_stress_coupling_kernel(kernel_name="eff_stress_x", variable="disp_x", component=0,
                                                             biot_coefficient=biot_coeff)
    builder.add_porous_flow_effective_stress_coupling_kernel(kernel_name="eff_stress_y", variable="disp_y", component=1,
                                                             biot_coefficient=biot_coeff)
    # builder.add_porous_flow_mass_volumetric_expansion_kernel(kernel_name="mass_exp", variable="pp") # <- Removed if we need one way coupling.

    fluid_property = SimpleFluidPropertiesConfig(name="water", bulk_modulus=2.2E9, viscosity=1.0E-3, density0=1000.0)
    builder.add_fluid_properties_config(fluid_property)
    builder.add_poromechanics_materials(
        fluid_properties_name="water",
        biot_coefficient=biot_coeff,
        solid_bulk_compliance=2E-11
    )


    builder.add_piecewise_function_from_data1d(name="injection_pressure_func", source_data1d=gauge_data_for_moose)

    builder.set_hydraulic_fracturing_bcs(
        injection_well_boundary_name="injection",
        injection_pressure_function_name="injection_pressure_func",
        confine_disp_x_boundaries="left right",
        confine_disp_y_boundaries="top bottom"
    )

    builder.add_standard_tensor_aux_vars_and_kernels({"stress": "stress", "total_strain": "strain"})

    # Add post-processors to model builder
    # This part is from baseline_model_builder.py (v1) which provides better post-processing options.
    shift_val_ft = kwargs.get('monitoring_point_shift_ft', 80)

    # Center point sampler, pressure
    builder.add_postprocessor(
        PointValueSamplerConfig(
            name="hf_center_pressure_sampler",
            variable="pp",
            point=(center_x_val, frac_y_center, 0)
        )
    )

    # Center point sampler, strain_yy
    builder.add_postprocessor(
        PointValueSamplerConfig(
            name="hf_center_strain_yy_sampler",
            variable="strain_yy",
            point=(center_x_val, frac_y_center, 0)
        )
    )

    # Monitoring point sampler, pressure
    builder.add_postprocessor(
        PointValueSamplerConfig(
            name="monitor_point_pressure_sampler",
            variable="pp",
            point=(center_x_val + shift_val_ft * conversion_factor, frac_y_center, 0)
        )
    )

    # Monitoring point sampler, strain_yy
    builder.add_postprocessor(
        PointValueSamplerConfig(
            name="monitor_point_strain_yy_sampler",
            variable="strain_yy",
            point=(center_x_val + shift_val_ft * conversion_factor, frac_y_center, 0)
        )
    )

    # Line sampler along the fracture, pressure
    builder.add_postprocessor(
        LineValueSamplerConfig(
            name="fiber_pressure_sampler",
            variable="pp",
            start_point=(center_x_val + shift_val_ft * conversion_factor,
                         domain_bounds[0] + kwargs.get("start_offset_y", 20) * conversion_factor, 0),
            end_point=(center_x_val + shift_val_ft * conversion_factor,
                       domain_bounds[1] - kwargs.get("end_offset_y", 20) * conversion_factor, 0),
            num_points=kwargs.get("num_fiber_points", 200),
            other_params={'sort_by': 'y'}
        )
    )

    # Line sampler along the fracture, strain_yy
    builder.add_postprocessor(
        LineValueSamplerConfig(
            name="fiber_strain_yy_sampler",
            variable="strain_yy",
            start_point=(center_x_val + shift_val_ft * conversion_factor,
                         domain_bounds[0] + kwargs.get("start_offset_y", 20) * conversion_factor, 0),
            end_point=(center_x_val + shift_val_ft * conversion_factor,
                       domain_bounds[1] - kwargs.get("end_offset_y", 20) * conversion_factor, 0),
            num_points=kwargs.get("num_fiber_points", 200),
            other_params={'sort_by': 'y'}
        )
    )

    # Time sequence stepper
    total_time = gauge_data_for_moose.taxis[-1] - gauge_data_for_moose.taxis[0]
    # Down sample two dataframes to speed up the simulation.
    timestepper_profile = gauge_data_for_moose.copy()

    dt_control_func = TimeSequenceStepper()
    dt_control_func.from_data1d(timestepper_profile)

    # Define the time stepper block
    builder.add_executioner_block(
        end_time=total_time,
        dt=3600 * 24 * 5,
        time_stepper_type='TimeSequenceStepper',
        stepper_config=dt_control_func
    )

    builder.add_initial_conditions_from_configs()
    builder.add_preconditioning_block(active_preconditioner='mumps')
    builder.add_outputs_block(exodus=False, csv=True, exodus_execute_on='FINAL')

    # Plot gauge data for reference
    fig, ax = plt.subplots()
    gauge_data_for_moose.plot(ax=ax, use_timestamp=True)
    plt.show()

    return builder

def post_processor_info_extractor(**kwargs) -> List[Data2D]:
    """
    This function extracts post-processor information from the simulation results.
    Will return two Data2D objects: one for pressure, one for strain_yy.

    :param output_dir: The directory containing the MOOSE output CSV files.
    :return: A list of Data2D objects representing the extracted post-processor data.
    """
    output_dir = kwargs.get("output_dir")
    if not output_dir:
        raise ValueError("output_dir must be provided in kwargs")

    vector_reader = MOOSEVectorPostProcessorReader()
    max_processor_id, _ = vector_reader.get_max_indices(output_dir)

    pressure_data2d = None
    strain_data2d = None

    for i in range(max_processor_id + 1):
        vector_reader.read(directory=output_dir, post_processor_id=i, variable_index=1)

        if "fiber_pressure_sampler" in vector_reader.sampler_name:
            pressure_data2d = vector_reader.to_analyzer()
        elif "fiber_strain_yy_sampler" in vector_reader.sampler_name:
            strain_data2d = vector_reader.to_analyzer()

    if pressure_data2d is None:
        raise FileNotFoundError("Could not find and extract 'fiber_pressure_sampler' data.")
    if strain_data2d is None:
        raise FileNotFoundError("Could not find and extract 'fiber_strain_yy_sampler' data.")

    # Post-processing:
    # Remove the first value for it is nan
    strain_data2d.data = strain_data2d.data[:, 1:]
    strain_data2d.taxis = strain_data2d.taxis[1:] - strain_data2d.taxis[1]

    pressure_data2d.data = pressure_data2d.data[:, 1:]
    pressure_data2d.taxis = pressure_data2d.taxis[1:] - pressure_data2d.taxis[1]

    # Only keep the \Delta P
    pressure_data2d.data = pressure_data2d.data - pressure_data2d.data[0, :]
    pressure_data2d.history.add_record("Pressure data adjusted to delta P by subtracting initial pressure profile.")

    return [pressure_data2d, strain_data2d]

def misfit_calculator(**kwargs) -> float:
    """
    This function calculates the misfit between simulated and observed data.

    :param kwargs: The placeholder for future parameters.
    :return: the calculated misfit value under the current configuration.
    """
    # Parameter set.
    conversion_factor = 0.3048  # feet to meters
    weight_matrix = kwargs.get('weight_matrix', np.array([1])) # The weight matrix must be a single value, or
    # a 1D array with uneven length. The value in the center must be the "center" of HF location.
    sim_fracture_center_ind: int = kwargs.get('sim_fracture_center_ind', 0) # The index of the fracture center in
    # the simulated data.
    observed_data_fracture_center_ind: int = kwargs.get('observed_data_fracture_center_ind', 0) # The index of
    # the fracture center in the observed data
    simulated_data: Data2D = kwargs.get('simulated_data')  # The simulated Data2D object
    observed_data: Data2D = kwargs.get('observed_data')  # The observed
    save_path = kwargs.get('save_path', None)
    simulated_data.daxis = simulated_data.daxis / conversion_factor  # Convert to feet

    if simulated_data is None or observed_data is None:
        raise ValueError("Both simulated_data and observed_data must be provided in kwargs")

    # Calculate misfit one channel by one channel.
    total_misfit = 0.0
    start_obs_ind = observed_data_fracture_center_ind - len(weight_matrix) // 2
    spacing_sim = simulated_data.daxis[1] - simulated_data.daxis[0]
    spacing_obs = observed_data.daxis[1] - observed_data.daxis[0]

    for iter_chan in range(len(weight_matrix)):
        # Extract the data points in the observed data
        obs_ind = start_obs_ind + iter_chan
        obs_depth = observed_data.daxis[obs_ind] # The channel we are calculating the misfit for.
        # Find the corresponding index in the simulated data
        # 1. get the depth value, which can be calculated from sim_fracture_center_ind and spacing
        spacing_chan_frac_obs = obs_depth - observed_data.daxis[observed_data_fracture_center_ind] # Depth difference
        sim_depth = simulated_data.daxis[sim_fracture_center_ind] + spacing_chan_frac_obs
        # 2. extract the data from simulated data
        # Init the data channel from simulated data
        sim_data_chan = simulated_data.get_value_by_depth(sim_depth)
        sim_data_chan_dataframe = Data1DGauge()
        sim_data_chan_dataframe.data = sim_data_chan
        sim_data_chan_dataframe.taxis = simulated_data.taxis
        sim_data_chan_dataframe.start_time = simulated_data.start_time
        sim_data_chan_dataframe.name = "simulated_channel"

        # 3. extract the data from observed data
        obs_data_chan = observed_data.get_value_by_depth(obs_depth)
        obs_data_chan_dataframe = Data1DGauge()
        obs_data_chan_dataframe.data = obs_data_chan
        obs_data_chan_dataframe.taxis = observed_data.taxis
        obs_data_chan_dataframe.start_time = observed_data.start_time
        obs_data_chan_dataframe.name = "observed_channel"

        # 4. calculate misfit for these two channels
        # Calibrate the two dataframes to have the same time axis
        sim_data_chan_dataframe.interpolate(new_taxis = obs_data_chan_dataframe.taxis,
                                            new_start_time = obs_data_chan_dataframe.start_time)

        # Initial time calibration
        sim_data_chan_dataframe.data -= sim_data_chan_dataframe.data[1]

        # Calculate misfit
        channel_misfit = np.sum((sim_data_chan_dataframe.data[1:-1] - obs_data_chan_dataframe.data[1:-1]) ** 2)
        # Weight the misfit
        weighted_channel_misfit = weight_matrix[iter_chan] * channel_misfit
        total_misfit += weighted_channel_misfit
        # For debugging
        print(f"Channel {iter_chan}: Depth {obs_depth:.2f} ft, "
              f"Simulated center index {sim_fracture_center_ind}, Observed center index {observed_data_fracture_center_ind}, "
              f"Channel misfit {channel_misfit:.4e}, Weighted channel misfit {weighted_channel_misfit:.4e}")

        # Create and save the plot
        fig, ax = plt.subplots()
        ax.plot(obs_data_chan_dataframe.taxis, obs_data_chan_dataframe.data, label='Observed')
        ax.plot(sim_data_chan_dataframe.taxis, sim_data_chan_dataframe.data, label='Simulated')
        ax.legend()
        ax.set_title(f'Channel {iter_chan} at Depth {obs_depth:.2f} ft')
        ax.set_xlabel('Time (s)')
        ax.set_ylabel('Value')
        misfit_text = f'Channel Misfit: {channel_misfit:.4e}\nWeighted Misfit: {weighted_channel_misfit:.4e}'
        ax.text(0.05, 0.95, misfit_text, transform=ax.transAxes, fontsize=10,
                verticalalignment='top', bbox=dict(boxstyle='round,pad=0.5', fc='yellow', alpha=0.5))

        if save_path:
            if not os.path.isdir(save_path):
                os.makedirs(save_path, exist_ok=True)
            plot_filename = os.path.join(save_path, f'misfit_channel_{iter_chan}_depth_{obs_depth:.2f}.png')
        else:
            plot_filename = f'misfit_channel_{iter_chan}_depth_{obs_depth:.2f}.png'

        plt.savefig(plot_filename)
        plt.close(fig)

    return total_misfit

if __name__ == "__main__":
    print("Don't run this script directly. It is meant to be imported as a module.\n--Shenyao")