import numpy as np
import random
import statistics

def mortality_probability(age, max_age=200):
    """
    A toy function that returns the probability of death at a given age.
    The function is designed so that average lifespan hovers near 200,
    but you can adjust it as needed.
    """
    # Example: A small base rate, increasing with age.
    # At age 0 --> prob ~ 0.0, at age 200 --> prob ~ 0.02
    # This will lead to many individuals reaching 200,
    # but is just a toy model. Tweak as necessary.
    return 0.0001 * age

def simulate_one_life(max_age=200):
    """
    Simulate a single life path, returning:
      - age of death (or 200 if still alive)
      - record of 'career changes'
      - record of 'major illnesses'
    """
    current_age = 0
    is_alive = True

    career_change_years = []
    illness_years = []

    while is_alive and current_age <= max_age:
        # Roll for mortality this year
        death_chance = mortality_probability(current_age, max_age)
        if random.random() < death_chance:
            # Person dies this year
            return current_age, career_change_years, illness_years

        # If still alive, check for events:
        # 1) Career change (example ages 25-150, 5% chance each year)
        if 25 <= current_age <= 150:
            if random.random() < 0.05:
                career_change_years.append(current_age)

        # 2) Major illness (example ages 40-190, 2% chance each year)
        if 40 <= current_age <= 190:
            if random.random() < 0.02:
                illness_years.append(current_age)

        current_age += 1

    # If we exited the loop without dying, it means we reached max_age
    return max_age, career_change_years, illness_years


def run_simulation(num_simulations=10000, max_age=200):
    """
    Run multiple simulations and return:
     - distribution of death ages
     - average age of death
     - average number of career changes
     - average number of major illnesses
    """
    death_ages = []
    career_change_counts = []
    illness_counts = []

    for _ in range(num_simulations):
        death_age, career_changes, illnesses = simulate_one_life(max_age)
        death_ages.append(death_age)
        career_change_counts.append(len(career_changes))
        illness_counts.append(len(illnesses))

    avg_death_age = statistics.mean(death_ages)
    avg_career_changes = statistics.mean(career_change_counts)
    avg_illnesses = statistics.mean(illness_counts)

    results = {
        "avg_death_age": avg_death_age,
        "avg_career_changes": avg_career_changes,
        "avg_illnesses": avg_illnesses,
        "death_age_distribution": death_ages
    }
    return results

# -----------------------
# Run the simulation
# -----------------------
if __name__ == "__main__":
    np.random.seed(42)  # For reproducibility (optional)
    random.seed(42)

    N = 10_000
    simulation_results = run_simulation(num_simulations=N, max_age=200)

    print(f"Ran {N} simulations of a 200-year lifespan model.\n")
    print(f"Average age of death: {simulation_results['avg_death_age']:.2f}")
    print(f"Average number of career changes: {simulation_results['avg_career_changes']:.2f}")
    print(f"Average number of major illnesses: {simulation_results['avg_illnesses']:.2f}")
    # You could also plot or further analyze death_age_distribution if desired.

import numpy as np
import random
import statistics
import plotly.graph_objects as go
from collections import Counter

# ----------------------
# Step 1: Define the simulation
# ----------------------
def mortality_probability(age, max_age=200):
    return 0.0001 * age

def simulate_one_life(max_age=200):
    current_age = 0
    career_change_years = []
    illness_years = []

    while current_age <= max_age:
        death_chance = mortality_probability(current_age, max_age)
        if random.random() < death_chance:
            return current_age, career_change_years, illness_years

        if 25 <= current_age <= 150:
            if random.random() < 0.05:
                career_change_years.append(current_age)
        if 40 <= current_age <= 190:
            if random.random() < 0.02:
                illness_years.append(current_age)

        current_age += 1
    return max_age, career_change_years, illness_years

def run_simulation(num_simulations=10000, max_age=200):
    death_ages = []
    career_change_counts = []
    illness_counts = []

    for _ in range(num_simulations):
        d_age, c_changes, i_ills = simulate_one_life(max_age)
        death_ages.append(d_age)
        career_change_counts.append(len(c_changes))
        illness_counts.append(len(i_ills))

    results = {
        "death_ages": death_ages,
        "career_change_counts": career_change_counts,
        "illness_counts": illness_counts
    }
    return results

# ----------------------
# Step 2: Binning functions
# ----------------------
def bin_career_changes(count):
    if count == 0:
        return "Career=0"
    elif count == 1:
        return "Career=1"
    elif count == 2:
        return "Career=2"
    else:
        return "Career=3+"

def bin_illnesses(count):
    if count == 0:
        return "Ill=0"
    elif count == 1:
        return "Ill=1"
    elif count == 2:
        return "Ill=2"
    else:
        return "Ill=3+"

def bin_death_age(age):
    if age <= 50:
        return "Age=0-50"
    elif age <= 100:
        return "Age=51-100"
    elif age <= 150:
        return "Age=101-150"
    else:
        return "Age=151-200"

# ----------------------
# Step 3: Build flow counters
# ----------------------
def build_flow_counts(results):
    career_counts = results["career_change_counts"]
    illness_counts = results["illness_counts"]
    death_ages = results["death_ages"]

    from collections import Counter
    combos = []

    # Build the list of (career_bin, ill_bin, age_bin)
    for c_count, i_count, d_age in zip(career_counts, illness_counts, death_ages):
        c_bin = bin_career_changes(c_count)
        i_bin = bin_illnesses(i_count)
        a_bin = bin_death_age(d_age)
        combos.append((c_bin, i_bin, a_bin))

    combo_counter = Counter(combos)

    flow1_counter = Counter()  # (career_bin -> ill_bin)
    flow2_counter = Counter()  # (ill_bin -> age_bin)

    for (c_bin, i_bin, a_bin), count in combo_counter.items():
        flow1_counter[(c_bin, i_bin)] += count
        flow2_counter[(i_bin, a_bin)] += count

    return flow1_counter, flow2_counter

# ----------------------
# Step 4: Convert counters -> Sankey format
# ----------------------
def build_sankey_data(flow1_counter, flow2_counter):
    career_bins = ["Career=0", "Career=1", "Career=2", "Career=3+"]
    illness_bins = ["Ill=0", "Ill=1", "Ill=2", "Ill=3+"]
    age_bins = ["Age=0-50", "Age=51-100", "Age=101-150", "Age=151-200"]

    all_nodes = career_bins + illness_bins + age_bins
    label_to_index = {label: i for i, label in enumerate(all_nodes)}

    source_list = []
    target_list = []
    value_list = []

    # (career -> illness)
    for (c_bin, i_bin), val in flow1_counter.items():
        source_list.append(label_to_index[c_bin])
        target_list.append(label_to_index[i_bin])
        value_list.append(val)

    # (illness -> age)
    for (i_bin, a_bin), val in flow2_counter.items():
        source_list.append(label_to_index[i_bin])
        target_list.append(label_to_index[a_bin])
        value_list.append(val)

    return all_nodes, source_list, target_list, value_list

# ----------------------
# Step 5: Plot with Plotly
# ----------------------
def plot_sankey(all_nodes, source_list, target_list, value_list, title="Life Simulation Sankey"):
    fig = go.Figure(data=[go.Sankey(
        node=dict(
            pad=15,
            thickness=20,
            line=dict(color="black", width=0.5),
            label=all_nodes
        ),
        link=dict(
            source=source_list,
            target=target_list,
            value=value_list
        )
    )])
    fig.update_layout(title_text=title, font_size=12)
    fig.show()

# ----------------------
# Main execution
# ----------------------
if __name__ == "__main__":
    # For reproducibility (optional)
    np.random.seed(42)
    random.seed(42)

    # 1) Run the simulation
    results = run_simulation(num_simulations=10000, max_age=200)

    # 2) Build the flow counters (career->illness, illness->death_age)
    flow1_counter, flow2_counter = build_flow_counts(results)

    # 3) Convert to Sankey data
    all_nodes, source_list, target_list, value_list = build_sankey_data(flow1_counter, flow2_counter)

    # 4) Plot Sankey
    plot_sankey(all_nodes, source_list, target_list, value_list,
                title="Monte Carlo Life Simulation (200-year span) - Sankey Diagram")