from dash import Dash, Input, Output, State, no_update
from data.functions import *
from transportEqs import *
from Kps_function.read_Kps_model import *
import os
from figure import build_figure
from main_layout import build_main_layout
from instructions import build_instructions_layout
from rst_info import build_rst_layout
from dash import html, dcc
import base64
import io

app = Dash(
    __name__,
    assets_folder=os.path.join(os.path.dirname(__file__), "assets"),
    suppress_callback_exceptions=True,
    #external_stylesheets=[dbc.themes.BOOTSTRAP],
)

# --------------------------------------------------------------------
# Layout
# --------------------------------------------------------------------
app.layout = html.Div([
    dcc.Location(id="url", refresh=False),

    html.Header(
        html.Div([
            html.A(
                html.Img(
                    src="/assets/gov-fda-new-white.svg",
                    alt="U.S. Food and Drug Administration",
                    className="fda-logo",
                ),
                href="https://www.fda.gov",
                target="_blank",
                className="site-brand",
            ),

            html.Span(
                "Extraction condition justification tool",
                className="site-tool-title",
            ),

            html.Nav([
                dcc.Link("Model", href="/"),
                dcc.Link("Quick start", href="/instructions"),
                dcc.Link("About this tool", href="/rst"),
            ], className="site-nav", **{"aria-label": "Main navigation"}),
        ], className="site-container site-header-inner"),
        className="site-header",
    ),

    html.Main(
        html.Div(id="page-content"),
        id="main-content",
        className="site-container site-main",
    ),

    html.Footer(
        html.Div("Office of Science and Engineering Laboratories · Contact", className="site-container"),
        className="site-footer",
    ),
])

app.validation_layout = html.Div([
    app.layout,
    build_main_layout(),
    build_instructions_layout(),
    build_rst_layout(),
])

@app.callback(
    Output("page-content", "children"),
    Input("url", "pathname")
)
def route(pathname):
    if pathname == "/instructions":
        return build_instructions_layout()
    elif pathname == "/rst":
        return build_rst_layout()
    else:
        return build_main_layout()

# --------------------------------------------------------------------
# Callbacks
# --------------------------------------------------------------------

@app.callback(
    Output("solute_upload_container", "style"),
    Input("solute_source", "value"),
)
def show_solute_upload(solute_source):

    if solute_source == "upload":
        return {"display": "block"}

    return {"display": "none"}

@app.callback(
    Output("uploaded_solute_data", "data"),
    Output("solute_upload_status", "children"),
    Input("solute_upload", "contents"),
    State("solute_upload", "filename"),
    prevent_initial_call=True,
)
def load_solute_file(contents, filename):

    if contents is None:
        return None, ""

    try:
        content_type, content_string = contents.split(",")

        decoded = base64.b64decode(content_string)

        uploaded_df = pd.read_excel(
            io.BytesIO(decoded)
        )

        required_columns = [
            "CASRN",
            "MW",
            "LogP",
            #"Vabc",
        ]

        missing = [
            col for col in required_columns
            if col not in uploaded_df.columns
        ]

        if missing:
            return (
                None,
                f"Missing required columns: {', '.join(missing)}"
            )

        numeric_columns = [
            "MW",
            "LogP",
            #"Vabc",
        ]

        for col in numeric_columns:
            uploaded_df[col] = pd.to_numeric(
                uploaded_df[col],
                errors="coerce"
            )

        if uploaded_df[numeric_columns].isna().any().any():
            return (
                None,
                "One or more required numeric values are missing or invalid."
            )

        return (
            uploaded_df.to_dict("records"),
            f"Loaded {len(uploaded_df)} solutes from {filename}"
        )

    except Exception as e:
        return None, f"Could not read file: {e}"

@app.callback(
    Output("c1_solvent", "disabled"),
    Output("c1_volume", "disabled"),
    Output("c1_swelling", "disabled"),
    Output("c1_temp", "disabled"),
    Output("c1_iterations", "disabled"),
    Input("condition1_mode", "value"),
    State("c1_solvent", "value"),
)
def toggle_condition1_inputs(mode, current_value):
    disable = (mode != "in_vitro")

    return disable, disable, disable, disable, disable

@app.callback(
    Output("c2_solvent", "disabled"),
    Output("c2_volume", "disabled"),
    Output("c2_swelling", "disabled"),
    Output("c2_temp", "disabled"),
    Output("c2_iterations", "disabled"),
    Input("condition2_mode", "value"),
    State("c2_solvent", "value"),
)
def toggle_condition2_inputs(mode, current_value):
    disable = (mode != "in_vitro")

    return disable, disable, disable, disable, disable

@app.callback(
    Output("mm0_chart", "figure"),
    Output("solute_calculate_warning", "children"),

    Input("calculate", "n_clicks"),
    State("matrix", "value"),
    State("tg", "value"),
    State("crystal", "value"),
    State("density", "value"),
    State("polymer_volume", "value"),
    State("surface_area", "value"),

    # Condition 1
    State("condition1_mode", "value"),
    State("c1_solvent", "value"),
    State("c1_volume", "value"),
    State("c1_swelling", "value"),
    State("c1_temp", "value"),
    State("c1_time", "value"),
    State("c1_iterations", "value"),
    # Condition 2
    State("condition2_mode", "value"),
    State("c2_solvent", "value"),
    State("c2_volume", "value"),
    State("c2_swelling", "value"),
    State("c2_temp", "value"),
    State("c2_time", "value"),
    State("c2_iterations", "value"),
    # MC
    State("samples", "value"),
    State("mc_option", "value"),
    # Solutes
    State("solute_source", "value"),
    State("uploaded_solute_data", "data"),
    prevent_initial_call=True,
)

def update_chart(n_clicks,
                 matrix, tg, crystal, density, polymer_volume, surface_area,
                 c1_mode, c1_solvent, c1_volume, c1_swelling,
                 c1_temp, c1_time, c1_iter,
                 c2_mode, c2_solvent, c2_volume,
                 c2_swelling, c2_temp, c2_time, c2_iter,samples,mc_option,solute_source,uploaded_solute_data):

    # ---------------------------------------------------------
    # Validate solute source
    # ---------------------------------------------------------
    if solute_source == "default":

        current_soluteData = soluteData.copy()

    elif solute_source == "upload":

        if not uploaded_solute_data:
            return (
                no_update,
                "Please upload a valid Excel file before calculating."
            )

        current_soluteData = pd.DataFrame.from_records(
            uploaded_solute_data
        )

    else:
        return (
            no_update,
            "Please select either the default or user-defined chemical set."
        )

    df = pd.DataFrame()
    df['CASRN'] = current_soluteData['CASRN'].to_numpy()

    current_soluteData["LogP"] = current_soluteData["LogP"].clip(upper=10)
    current_soluteData["Vabc"] = current_soluteData["MW"]

    N = int(samples)
    pindex = np.where(polymers == matrix)[0][0]
    CHRIS_category = categories[pindex]
    Polymer_Density = float(density)
    Polymer_Tg = float(tg) + 273.15
    Polymer_X = float(crystal)/100.

    L = float(polymer_volume)/float(surface_area)

    # Condition 1
    ExtractionTime = float(c1_time) * 3600.
    if c1_mode == 'in_vivo_conservative':
        c1_dist = ConservativeMonteCarlo(current_soluteData, CHRIS_category, L, ExtractionTime, N)

    elif c1_mode == 'in_vivo_tissue':
        c1_dist = TissueMonteCarlo(current_soluteData, CHRIS_category, L, ExtractionTime, N)

    else:
        Solvent_Name = c1_solvent
        Solvent_MW = Solvent_MWs[Solvent_Name]
        Solvent_Density = Solvent_Densities[Solvent_Name]
        Solvent_PI = Solvent_PIs[Solvent_Name]
        ExtractionT = float(c1_temp) + 273.15
        Ms_M0 = float(c1_swelling)
        Swell = 1.+(Ms_M0-1.)*Polymer_Density/Solvent_Density
        w = (Ms_M0 - 1.)/(Ms_M0-Polymer_X)
        Iterations = float(c1_iter)
        Lbath = float(c1_volume) / float(surface_area)

        c1_dist = ExtractMonteCarlo(current_soluteData, Solvent_PI, w, ExtractionT, Polymer_Tg, Solvent_Name, Solvent_MW,
                                                    ExtractionTime, Swell, Iterations, CHRIS_category, L, Lbath, N)

    # Condition 2
    ExtractionTime = float(c2_time) * 3600.
    if c2_mode == 'in_vivo_conservative':
        c2_dist = ConservativeMonteCarlo(current_soluteData, CHRIS_category, L, ExtractionTime, N)

    elif c2_mode == 'in_vivo_tissue':
        c2_dist = TissueMonteCarlo(current_soluteData, CHRIS_category, L, ExtractionTime, N)

    else:
        Solvent_Name = c2_solvent
        Solvent_MW = Solvent_MWs[Solvent_Name]
        Solvent_Density = Solvent_Densities[Solvent_Name]
        Solvent_PI = Solvent_PIs[Solvent_Name]
        ExtractionT = float(c2_temp) + 273.15
        Ms_M0 = float(c2_swelling)
        Swell = 1. + (Ms_M0 - 1.) * Polymer_Density / Solvent_Density
        w = (Ms_M0 - 1.) / (Ms_M0 - Polymer_X)
        Iterations = float(c2_iter)
        Lbath = float(c2_volume) / float(surface_area)

        c2_dist = ExtractMonteCarlo(current_soluteData, Solvent_PI, w, ExtractionT, Polymer_Tg, Solvent_Name, Solvent_MW,
                                                    ExtractionTime, Swell, Iterations, CHRIS_category, L, Lbath, N)

    if mc_option == "absolute":
        c1_lower, c1_median, c1_upper = np.percentile(
            c1_dist,
            [25, 50, 75],
            axis=1
        )

        c2_lower, c2_median, c2_upper = np.percentile(
            c2_dist,
            [25, 50, 75],
            axis=1
        )

        df['Condition 1'] = c1_median
        df['Cond1_err_plus'] = c1_upper - c1_median
        df['Cond1_err_minus'] = c1_median - c1_lower

        df['Condition 2'] = c2_median
        df['Cond2_err_plus'] = c2_upper - c2_median
        df['Cond2_err_minus'] = c2_median - c2_lower

    elif mc_option == "ratio":

        ratio_dist = np.divide(
            c1_dist,
            c2_dist,
            out=np.full_like(c1_dist, np.nan),
            where=c2_dist > 0
        )

        ratio_lower, ratio_median, ratio_upper = np.percentile(
            ratio_dist,
            [25, 50, 75],
            axis=1
        )

        df['Ratio'] = ratio_median
        df['Ratio_err_plus'] = ratio_upper - ratio_median
        df['Ratio_err_minus'] = ratio_median - ratio_lower

    fig = build_figure(df, mc_option)

    return fig, ""

def ConservativeMonteCarlo(soluteData, CHRIS_category, L, time, N):

    nSolutes = len(soluteData)
    distributions = zeros((nSolutes, N))

    for iSolute in range(nSolutes):

        Solute_MW = soluteData['MW'].iloc[iSolute]

        D_CHRIS = get_D_CHRIS(
            Solute_MW,
            CHRIS_category,
            N=N
        )

        tau = D_CHRIS * time / L**2
        mass = Conservative(tau)

        distributions[iSolute, :] = mass

    return distributions

def TissueMonteCarlo(soluteData, CHRIS_category, L, time, N):

    nSolutes = len(soluteData)
    distributions = zeros((nSolutes, N))

    for iSolute in range(nSolutes):

        Solute_MW = soluteData['MW'].iloc[iSolute]
        Solute_logP = soluteData['LogP'].iloc[iSolute]

        D_CHRIS = get_D_CHRIS(
            Solute_MW,
            CHRIS_category,
            N=N
        )

        tau = D_CHRIS * time / L**2

        Dt = get_Dt(N)
        Kpt = get_Kpt(Solute_logP, N)

        beta = (1. / Kpt) * sqrt(Dt / D_CHRIS)

        mass = SolubilityLimited(beta, tau)

        distributions[iSolute, :] = mass

    return distributions

def ExtractMonteCarlo(
        soluteData, Solvent_PI, w, ExtractionT,
        Polymer_Tg, Solvent_Name, Solvent_MW,
        ExtractionTime, Swell, Iterations,
        CHRIS_category, L, Lbath, N):

    nSolutes = len(soluteData)
    distributions = zeros((nSolutes, N))

    for iSolute in range(nSolutes):

        Solute_MW = soluteData['MW'].iloc[iSolute]
        Solute_logP = soluteData['LogP'].iloc[iSolute]
        Solute_Vabc = soluteData['Vabc'].iloc[iSolute]

        if Solvent_Name == 'blood':

            Kps = get_Kpb(Solute_logP, N)

        else:

            ToKPS = array([[Solute_logP, Solvent_PI]])

            model = QuantileGridFromCoeffs(
                export_dir='Kps_function/Kps_model'
            )

            Kps = 10 ** model.sample(
                ToKPS,
                n_samples=N
            )[0]

        D_Extract = get_D_Extract(
            w,
            ExtractionT,
            Polymer_Tg,
            Solvent_Name,
            Solvent_MW,
            Solute_MW,
            Solute_Vabc,
            CHRIS_category,
            N=N
        )

        tau = D_Extract * ExtractionTime / L**2
        alpha = Lbath / L / Kps
        massExtraction = zeros(N)

        for i in range(N):
            massExtraction[i] = Extraction(
                tau[i],
                alpha[i],
                Kps[i],
                Swell,
                Iterations
            )

        distributions[iSolute, :] = massExtraction

    return distributions


if __name__ == "__main__":
    app.run(debug=True)