"""
Streamlit app: simulate voting scenarios for "Your Party" in Oxford East,
with additional static scenario charts (two univariate plots and one surface/heatmap)
showing outcomes as percent-left and percent-right switching values vary.

Usage:
    pip install streamlit plotly pandas numpy
    streamlit run this_file.py
"""

import streamlit as st
import pandas as pd
import numpy as np
import plotly.graph_objects as go

# --------------------------
# Fixed inputs (from user)
# --------------------------
E = 71845                        # Electorate
T = 54.8                         # Historic turnout percent
P_L = 31.1                       # Labour percent
P_D = 8.8                        # LibDem percent
P_C = 7.1                        # Conservative percent
P_R = 17.5                       # Reform percent
P_G = 24.0                       # Green percent

PARTY_COLORS = {
    "Labour": "#E4003B",
    "LibDem": "#FAA61A",
    "Conservative": "#0087DC",
    "Reform": "#7F3FBF",
    "Green": "#78AB46",
    "Your Party": "#FFFFFF",
    # for plotting aggregated "Right party" / "Other" states in scenario maps
    "RightWinner": "#0087DC",
    "Other": "#999999",
}

st.set_page_config(layout="wide", page_title="Oxford East: Your Party simulator")
st.title("Oxford East — Your Party voting simulator")

# --------------------------
# Derived baseline numbers
# --------------------------
historic_voters = E * (T / 100.0)
non_voters_pool = E - historic_voters

baseline = {
    "Labour": historic_voters * (P_L / 100.0),
    "LibDem": historic_voters * (P_D / 100.0),
    "Conservative": historic_voters * (P_C / 100.0),
    "Reform": historic_voters * (P_R / 100.0),
    "Green": historic_voters * (P_G / 100.0),
    "Your Party": 0.0,
}

# --------------------------
# User-controllable variables
# --------------------------
st.sidebar.header("Simulation inputs")

max_non_voters = int(non_voters_pool)
non_voters_convinced = st.sidebar.number_input(
    "Non-voters convinced to vote for Your Party (absolute number)",
    min_value=0,
    max_value=max_non_voters,
    value=0,
    step=1,
)

percent_left_switch = st.sidebar.slider(
    "Percent of left voters who switch to Your Party",
    min_value=0.0,
    max_value=100.0,
    value=0.0,
    step=0.1,
)

percent_right_switch = st.sidebar.slider(
    "Percent of right voters who switch to Your Party",
    min_value=0.0,
    max_value=100.0,
    value=0.0,
    step=0.1,
)

st.subheader("Variables in this simulation")
st.write(f"- Non-voters convinced to vote for Your Party: {non_voters_convinced:,}")
st.write(f"- Percent of left voters convinced (from left pool): {percent_left_switch:.1f}%")
st.write(f"- Percent of right voters convinced (from right pool): {percent_right_switch:.1f}%")

# --------------------------
# Helper: proportional reduction
# --------------------------
def proportional_reduction(baseline_votes, group_keys, transfers_amount):
    reduced = {}
    group_total = sum(baseline_votes[k] for k in group_keys)
    if group_total <= 0 or transfers_amount <= 0:
        for k in group_keys:
            reduced[k] = baseline_votes[k]
        return reduced
    for k in group_keys:
        share = baseline_votes[k] / group_total
        reduced[k] = baseline_votes[k] - transfers_amount * share
    return reduced

# --------------------------
# Function to compute final votes for any pair of (pct_left, pct_right),
# keeping non_voters_convinced fixed to the UI value.
# --------------------------
def compute_final_votes(pct_left, pct_right, non_voters_fixed):
    left_pool_votes = baseline["Labour"] + baseline["Green"]
    right_pool_votes = baseline["Conservative"] + baseline["Reform"]
    left_transfers = left_pool_votes * (pct_left / 100.0)
    right_transfers = right_pool_votes * (pct_right / 100.0)

    reduced_left = proportional_reduction(baseline, ["Labour", "Green"], left_transfers)
    reduced_right = proportional_reduction(baseline, ["Conservative", "Reform"], right_transfers)

    final = {
        "Labour": reduced_left["Labour"],
        "Green": reduced_left["Green"],
        "Conservative": reduced_right["Conservative"],
        "Reform": reduced_right["Reform"],
        "LibDem": baseline["LibDem"],
        "Your Party": baseline["Your Party"] + left_transfers + right_transfers + non_voters_fixed,
    }
    return final

# --------------------------
# Current single-scenario outcome (keeps behavior of the previous app)
# --------------------------
final_votes = compute_final_votes(percent_left_switch, percent_right_switch, non_voters_convinced)
df = pd.DataFrame({"party": list(final_votes.keys()), "votes": list(final_votes.values())})
df = df.sort_values("votes", ascending=False).reset_index(drop=True)
top_party = df.loc[0, "party"]

# Plot horizontal bars (current scenario)
bar_colors = [PARTY_COLORS.get(p, "#999999") for p in df["party"]]
marker_line_width = [4 if p == top_party else 1 for p in df["party"]]
marker_line_color = ["black" if p == top_party else "rgba(0,0,0,0.3)" for p in df["party"]]

fig_main = go.Figure()
fig_main.add_trace(
    go.Bar(
        x=df["votes"],
        y=df["party"],
        orientation="h",
        marker=dict(color=bar_colors, line=dict(color=marker_line_color, width=marker_line_width)),
        hovertemplate="%{y}: %{x:,.0f} votes<extra></extra>",
    )
)
fig_main.update_layout(xaxis_title="Votes (absolute number)", yaxis_title="Party",
                       margin=dict(l=140, r=30, t=60, b=40), height=420, showlegend=False)
fig_main.add_annotation(
    x=df.loc[0, "votes"], y=df.loc[0, "party"], xanchor="left", yanchor="middle",
    text=f" Top: {df.loc[0,'party']} ({df.loc[0,'votes']:.0f} votes)",
    font=dict(size=12, color="black", family="Arial"),
    showarrow=False, bgcolor="rgba(255,255,255,0.8)", bordercolor="black", borderwidth=1,
)
st.plotly_chart(fig_main, use_container_width=True)

col1, col2, col3 = st.columns(3)
new_total_votes = sum(final_votes.values())
new_turnout = new_total_votes / E * 100.0
col1.metric("Total votes (new)", f"{new_total_votes:,.0f}")
col2.metric("New turnout", f"{new_turnout:.1f}%")
col3.metric("Your Party votes", f"{final_votes['Your Party']:,.0f}")

with st.expander("Show numeric results"):
    st.dataframe(df.style.format({"votes": "{:,.0f}"}), height=260)

# --------------------------
# Static scenario charts requested:
# - Two univariate plots:
#     1) Vary percent_left (0..max_range) with percent_right fixed at current UI value.
#     2) Vary percent_right (0..max_range) with percent_left fixed at current UI value.
#   For each univariate plot show Your Party votes and mark ranges where Your Party is the winner.
# - One surface plot (heatmap) over grid of (pct_left, pct_right), showing winner.
#
# Also compute and state the ranges of (pct_left, pct_right) where:
# - Your Party wins
# - A right party wins (Conservative or Reform)
# --------------------------
st.markdown("---")
st.header("Scenario maps — how switching percentages change outcomes")

# grid resolution and ranges
max_range = 60  # percent range to explore for each axis (0..max_range)
step = 1
left_values = np.arange(0.0, max_range + step, step)
right_values = np.arange(0.0, max_range + step, step)

# Precompute grid outcomes (winner and Your Party votes)
grid_winner = np.empty((len(left_values), len(right_values)), dtype=np.int8)  # 1=Your Party, 2=Right party, 0=Other
grid_your_votes = np.empty_like(grid_winner, dtype=float)

for i, lp in enumerate(left_values):
    for j, rp in enumerate(right_values):
        fv = compute_final_votes(lp, rp, non_voters_convinced)
        # determine winner
        winner = max(fv.items(), key=lambda kv: kv[1])[0]
        grid_your_votes[i, j] = fv["Your Party"]
        if winner == "Your Party":
            grid_winner[i, j] = 1
        elif winner in ("Conservative", "Reform"):
            grid_winner[i, j] = 2
        else:
            grid_winner[i, j] = 0

# Determine where Your Party wins and where a right party wins
your_win_points = np.argwhere(grid_winner == 1)
right_win_points = np.argwhere(grid_winner == 2)

def describe_axis_ranges(points, axis_values):
    # returns min and max axis values where points exist, or None if empty
    if len(points) == 0:
        return None, None
    vals = axis_values[points]
    return float(vals.min()), float(vals.max())

# For a compact description, compute bounding box ranges (min/max left and right percents) for Your Party and right party
if your_win_points.size > 0:
    your_left_min = left_values[your_win_points[:, 0]].min()
    your_left_max = left_values[your_win_points[:, 0]].max()
    your_right_min = right_values[your_win_points[:, 1]].min()
    your_right_max = right_values[your_win_points[:, 1]].max()
else:
    your_left_min = your_left_max = your_right_min = your_right_max = None

if right_win_points.size > 0:
    r_left_min = left_values[right_win_points[:, 0]].min()
    r_left_max = left_values[right_win_points[:, 0]].max()
    r_right_min = right_values[right_win_points[:, 1]].min()
    r_right_max = right_values[right_win_points[:, 1]].max()
else:
    r_left_min = r_left_max = r_right_min = r_right_max = None

# Percent of grid where each outcome occurs
total_cells = grid_winner.size
your_win_pct = (grid_winner == 1).sum() / total_cells * 100.0
right_win_pct = (grid_winner == 2).sum() / total_cells * 100.0

# Print concise range summaries
st.subheader("Outcome ranges (within scanned 0–{}%)".format(max_range))
col_a, col_b = st.columns(2)
with col_a:
    if your_win_points.size > 0:
        st.write("**Your Party wins** for roughly:")
        st.write(f"- Left switching: {your_left_min:.0f}% – {your_left_max:.0f}%")
        st.write(f"- Right switching: {your_right_min:.0f}% – {your_right_max:.0f}%")
        st.write(f"- Grid coverage: {your_win_pct:.1f}% of scanned scenarios")
    else:
        st.write("**Your Party does not win** in any scanned scenario (0–{}%).".format(max_range))
with col_b:
    if right_win_points.size > 0:
        st.write("**A right party (Conservative or Reform) wins** for roughly:")
        st.write(f"- Left switching: {r_left_min:.0f}% – {r_left_max:.0f}%")
        st.write(f"- Right switching: {r_right_min:.0f}% – {r_right_max:.0f}%")
        st.write(f"- Grid coverage: {right_win_pct:.1f}% of scanned scenarios")
    else:
        st.write("**Right parties do not win** in any scanned scenario (0–{}%).".format(max_range))

# --------------------------
# Univariate plot 1: vary left percent, keep right fixed at current UI value
# --------------------------
your_votes_left = []
winners_left = []
for lp in left_values:
    fv = compute_final_votes(lp, percent_right_switch, non_voters_convinced)
    your_votes_left.append(fv["Your Party"])
    winners_left.append(max(fv.items(), key=lambda kv: kv[1])[0])

fig_left = go.Figure()
fig_left.add_trace(go.Scatter(x=left_values, y=your_votes_left, mode="lines", name="Your Party votes",
                              hovertemplate="Left switch: %{x:.0f}%<br>Your Party votes: %{y:,.0f}<extra></extra>"))
# shade region where Your Party is the winner
is_your_winner = np.array([1 if w == "Your Party" else 0 for w in winners_left])
if is_your_winner.any():
    # create filled area under the line where your party wins
    win_x = left_values * is_your_winner
    win_y = np.array(your_votes_left) * is_your_winner
    # convert zeros to nan so fill only covers winning points
    win_x = np.where(is_your_winner, win_x, np.nan)
    win_y = np.where(is_your_winner, win_y, np.nan)
    fig_left.add_trace(go.Scatter(x=win_x, y=win_y, fill="tozeroy", name="Your Party wins region",
                                  hoverinfo="skip", opacity=0.25, marker=dict(color=PARTY_COLORS["Your Party"])))

fig_left.update_layout(title=f"Vary percent left switching (right fixed at {percent_right_switch:.1f}%)",
                       xaxis_title="Percent of left voters switching to Your Party",
                       yaxis_title="Your Party votes (absolute)",
                       height=350, margin=dict(l=80, r=20, t=50, b=40))
st.plotly_chart(fig_left, use_container_width=True)

# --------------------------
# Univariate plot 2: vary right percent, keep left fixed at current UI value
# --------------------------
your_votes_right = []
winners_right = []
for rp in right_values:
    fv = compute_final_votes(percent_left_switch, rp, non_voters_convinced)
    your_votes_right.append(fv["Your Party"])
    winners_right.append(max(fv.items(), key=lambda kv: kv[1])[0])

fig_right = go.Figure()
fig_right.add_trace(go.Scatter(x=right_values, y=your_votes_right, mode="lines", name="Your Party votes",
                               hovertemplate="Right switch: %{x:.0f}%<br>Your Party votes: %{y:,.0f}<extra></extra>"))
is_your_winner_r = np.array([1 if w == "Your Party" else 0 for w in winners_right])
if is_your_winner_r.any():
    win_x = right_values * is_your_winner_r
    win_y = np.array(your_votes_right) * is_your_winner_r
    win_x = np.where(is_your_winner_r, win_x, np.nan)
    win_y = np.where(is_your_winner_r, win_y, np.nan)
    fig_right.add_trace(go.Scatter(x=win_x, y=win_y, fill="tozeroy", name="Your Party wins region",
                                   hoverinfo="skip", opacity=0.25, marker=dict(color=PARTY_COLORS["Your Party"])))

fig_right.update_layout(title=f"Vary percent right switching (left fixed at {percent_left_switch:.1f}%)",
                        xaxis_title="Percent of right voters switching to Your Party",
                        yaxis_title="Your Party votes (absolute)",
                        height=350, margin=dict(l=80, r=20, t=50, b=40))
st.plotly_chart(fig_right, use_container_width=True)

# --------------------------
# Surface / heatmap plot showing winner across the grid
# Use integer mapping: 0=Other/Left/LibDem, 1=Your Party, 2=Right party
# --------------------------
z = grid_winner  # shape (len(left_values), len(right_values))

# Map winner codes to colors and labels
color_for_code = {
    0: PARTY_COLORS["Other"],
    1: PARTY_COLORS["Your Party"],
    2: PARTY_COLORS["RightWinner"],
}

# Build a simple three-stop colorscale: 0 -> code0 color, 0.5 -> code1 color, 1 -> code2 color
# z ranges from 0..2, so we normalize positions across [0,1]
colorscale = [
    [0.0, color_for_code[0]],
    [0.5, color_for_code[1]],
    [1.0, color_for_code[2]],
]

# Create heatmap: x axis = right percent, y axis = left percent
fig_heat = go.Figure(data=go.Heatmap(
    z=z,
    x=right_values,
    y=left_values,
    colorscale=colorscale,
    zmin=0,
    zmax=2,
    colorbar=dict(
        title="Winner",
        tickmode="array",
        tickvals=[0, 1, 2],
        ticktext=["Other", "Your Party", "Right party"],
    ),
    hovertemplate="Left %{y:.0f}%<br>Right %{x:.0f}%<br>Winner: %{z}<extra></extra>",
))
fig_heat.update_layout(
    title="Outcome surface (grid): winner by percent-left / percent-right",
    xaxis_title="Percent right switching to Your Party",
    yaxis_title="Percent left switching to Your Party",
    height=600,
    margin=dict(l=80, r=40, t=60, b=60),
)
st.plotly_chart(fig_heat, use_container_width=True)

# --------------------------
# Additional numeric summary: numerical ranges of interest as lists (not only bounding boxes)
# Provide examples of threshold pairs where Your Party first becomes top (approx)
# --------------------------
st.markdown("### Additional notes and diagnostics")
# Find minimal sum of switches at which Your Party wins: example diagnostics
your_indices = np.argwhere(z == 1)
if your_indices.size > 0:
    # compute a "distance" metric (left+right) and find minimal combined switching where Your Party wins
    sums = left_values[your_indices[:, 0]] + right_values[your_indices[:, 1]]
    idx_min = np.argmin(sums)
    lp_min = left_values[your_indices[idx_min, 0]]
    rp_min = right_values[your_indices[idx_min, 1]]
    st.write(f"- Earliest (by sum of switches) scenario where Your Party wins in scanned grid: "
             f"Left = {lp_min:.0f}%, Right = {rp_min:.0f}% (sum = {lp_min+rp_min:.0f}%).")
else:
    st.write("- Your Party never wins in the scanned 0–{}% grid.".format(max_range))

if right_win_points.size > 0:
    # similarly earliest right-party win by sum
    r_indices = np.argwhere(z == 2)
    sums_r = left_values[r_indices[:, 0]] + right_values[r_indices[:, 1]]
    idx_min_r = np.argmin(sums_r)
    lr_min = left_values[r_indices[idx_min_r, 0]]
    rr_min = right_values[r_indices[idx_min_r, 1]]
    st.write(f"- Earliest (by sum of switches) scenario where a right party wins: "
             f"Left = {lr_min:.0f}%, Right = {rr_min:.0f}% (sum = {lr_min+rr_min:.0f}%).")
else:
    st.write("- Right parties never win in the scanned 0–{}% grid.".format(max_range))

st.caption("Notes: grids and ranges are computed for percent values between 0 and {} (step {}). "
           "Non-voters convinced is held fixed to the value selected in the sidebar for these scenario maps."
           .format(max_range, step))