import gradio as gr
import pandas as pd
import numpy as np
from numpy.random import default_rng
import io  # For BytesIO to handle file in memory

# 1. Data Generation Function (customizable via UI filters)
def generate_custom_model_points(
    mp_count_val, seed_val, age_min_val, age_max_val,
    sa_min_val, sa_max_val, policy_terms_selection_val,
    include_sex_val, policy_count_fixed_val
):
    """
    Generates seriatim model points based on user-defined parameters.
    """
    rng = default_rng(int(seed_val))
    mp_count_val = int(mp_count_val)
    age_min_val = int(age_min_val)
    age_max_val = int(age_max_val)
    sa_min_val = float(sa_min_val)
    sa_max_val = float(sa_max_val)

    # Policy ID
    policy_id_col = np.arange(1, mp_count_val + 1)

    # Issue Age
    age_at_entry = rng.integers(low=age_min_val, high=age_max_val + 1, size=mp_count_val)

    # Sex
    if include_sex_val:
        sex_options = ["M", "F"]
        sex_col = np.fromiter(map(lambda i: sex_options[i], rng.integers(low=0, high=len(sex_options), size=mp_count_val)), np.dtype('<U1'))
    else:
        sex_col = np.full(mp_count_val, "U")  # Unknown/Unspecified if not included

    # Policy Term
    if not policy_terms_selection_val:  # Default if user deselects all
        policy_terms_selection_val = [10, 15, 20]
    policy_term_options = np.array(policy_terms_selection_val).astype(int)
    policy_term_col = rng.choice(policy_term_options, size=mp_count_val)

    # Sum Assured
    sum_assured_col = np.round((sa_max_val - sa_min_val) * rng.random(size=mp_count_val) + sa_min_val, -3)

    # Duration in month: 1 <= Duration(mth) < Policy Term in month
    # (policy_term_col * 12) is term in months.
    # Max duration is (term_in_months - 1). Min duration is 1.
    max_duration_val = policy_term_col * 12 - 1
    # Ensure max_duration_val is at least 1 to prevent issues with rng.random if term is very short
    max_duration_val = np.maximum(1, max_duration_val)
    duration_mth_col = (rng.random(size=mp_count_val) * max_duration_val).astype(int) + 1
    # Ensure duration_mth does not exceed max_duration_val (especially if max_duration_val was 1)
    duration_mth_col = np.minimum(duration_mth_col, max_duration_val)
    # And ensure it's at least 1
    duration_mth_col = np.maximum(1, duration_mth_col)


    # Policy Count
    if policy_count_fixed_val:
        policy_count_col_val = np.ones(mp_count_val, dtype=int)
    else:
        policy_count_col_val = rng.integers(low=1, high=101, size=mp_count_val)  # Variable 1-100

    # Create DataFrame with policy_id as the first column
    data_dict = {
        "policy_id": policy_id_col,
        "age_at_entry": age_at_entry,
        "sex": sex_col,
        "policy_term": policy_term_col,
        "policy_count": policy_count_col_val,
        "sum_assured": sum_assured_col,
        "duration_mth": duration_mth_col
    }

    model_point_df = pd.DataFrame(data_dict)

    return model_point_df

# 2. Gradio App Definition
with gr.Blocks() as demo:
    gr.Markdown("# Actuarial Model Points Generator")
    gr.Markdown(
        "Configure the parameters below to generate a custom set of seriatim model points. "
        "The generated table can be viewed and downloaded as an Excel file."
    )

    df_state = gr.State()  # To hold the generated DataFrame

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Generation Parameters")

            mp_count_input = gr.Slider(
                minimum=100, maximum=50000, value=10000, step=100,
                label="Number of Model Points"
            )
            seed_input = gr.Number(
                value=12345, precision=0, label="Random Seed"
            )

            gr.Markdown("#### Age Parameters")
            age_min_input = gr.Slider(
                minimum=18, maximum=40, value=20, step=1, label="Minimum Age at Entry"
            )
            age_max_input = gr.Slider(
                minimum=41, maximum=80, value=59, step=1, label="Maximum Age at Entry"
            )

            gr.Markdown("#### Sum Assured Parameters ($)")
            sum_assured_min_input = gr.Number(
                value=10000, label="Minimum Sum Assured"
            )
            sum_assured_max_input = gr.Number(
                value=1000000, label="Maximum Sum Assured"
            )

            gr.Markdown("#### Policy Options")
            policy_terms_input = gr.CheckboxGroup(
                choices=[5, 10, 15, 20, 25, 30], value=[10, 15, 20],
                label="Available Policy Terms (Years)"
            )
            include_sex_input = gr.Checkbox(
                value=True, label="Include Sex (M/F)"
            )
            policy_count_fixed_input = gr.Checkbox(
                value=True, label="Fixed Policy Count = 1 (Uncheck for variable count 1-100)"
            )

            generate_btn = gr.Button("Generate Model Points", variant="primary")

        with gr.Column(scale=2):
            model_points_display = gr.Dataframe(label="Generated Model Points")
            download_excel_btn = gr.DownloadButton(
                label="Download Excel",
                value="model_points.xlsx", # Default filename
                variant="secondary"
            )

    # 3. Event Handlers
    def handle_generate_button_click(
        mp_c, s, age_m, age_mx, sa_m, sa_mx, p_terms, incl_sex, pc_fixed
    ):
        if int(age_m) >= int(age_mx):
            gr.Warning("Minimum Age must be less than Maximum Age.")
            return df_state.value, df_state.value # Keep current table and state
        if float(sa_m) >= float(sa_mx):
            gr.Warning("Minimum Sum Assured must be less than Maximum Sum Assured.")
            return df_state.value, df_state.value
        if not p_terms:
            gr.Warning("No Policy Terms selected. Using defaults: [10, 15, 20].")
            # Generation function will handle default if p_terms is empty list

        gr.Info("Generating model points... Please wait.")
        df = generate_custom_model_points(
            mp_c, s, age_m, age_mx, sa_m, sa_mx, p_terms, incl_sex, pc_fixed
        )
        gr.Info(f"{len(df)} model points generated successfully!")
        return df, df

    def handle_download_button_click(current_df_to_download):
        if current_df_to_download is None or current_df_to_download.empty:
            gr.Warning("No data available to download. Generate model points first.")
            empty_excel_output = io.BytesIO()
            pd.DataFrame().to_excel(empty_excel_output, index=False)
            empty_excel_output.seek(0)
            return empty_excel_output

        excel_output = io.BytesIO()
        current_df_to_download.to_excel(excel_output, sheet_name='ModelPoints', engine='xlsxwriter', index=False)
        excel_output.seek(0)
        return excel_output

    # Wire up UI components to functions
    inputs_list = [
        mp_count_input, seed_input, age_min_input, age_max_input,
        sum_assured_min_input, sum_assured_max_input, policy_terms_input,
        include_sex_input, policy_count_fixed_input
    ]

    generate_btn.click(
        fn=handle_generate_button_click,
        inputs=inputs_list,
        outputs=[model_points_display, df_state]
    )

    download_excel_btn.click(
        fn=handle_download_button_click,
        inputs=[df_state],
        outputs=[download_excel_btn]
    )


if __name__ == "__main__":
    demo.launch()