# pyrefly: ignore [missing-import]
import streamlit as st
import numpy as np
import pandas as pd
from scipy.optimize import linprog
import plotly.graph_objects as go
import plotly.express as px
import json
from db_utils import get_connection, hash_password, verify_password, derive_key, encrypt_data, decrypt_data

# Set page configurations
st.set_page_config(
    page_title="Marketing Budget Optimization System",
    page_icon="📊",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ---------- User Authentication ----------
if 'logged_in' not in st.session_state:
    st.session_state['logged_in'] = False
    st.session_state['user_id'] = None
    st.session_state['username'] = ''
    st.session_state['password'] = ''

def login_user(username, password):
    conn = get_connection()
    cur = conn.cursor()
    cur.execute("SELECT id, password_hash, salt FROM users WHERE username = ?", (username,))
    row = cur.fetchone()
    conn.close()
    if row and verify_password(row['password_hash'], row['salt'], password):
        st.session_state['logged_in'] = True
        st.session_state['user_id'] = row['id']
        st.session_state['username'] = username
        st.session_state['password'] = password
        st.success(f"Logged in as {username}")
    else:
        st.error("Invalid username or password")

def register_user(username, password, full_name, email, company):
    pwd_hash, salt = hash_password(password)
    conn = get_connection()
    cur = conn.cursor()
    try:
        cur.execute("INSERT INTO users (username, password_hash, salt) VALUES (?,?,?)", (username, pwd_hash, salt))
        conn.commit()
        st.success("Registration successful. You can now log in.")
    except Exception as e:
        st.error(f"Registration failed: {e}")
    finally:
        conn.close()

# Sidebar Account Section
with st.sidebar.expander("User Account", expanded=True):
    if not st.session_state['logged_in']:
        mode = st.radio("Choose action", ["Login", "Register"])
        if mode == "Login":
            login_user_name = st.text_input("Username")
            login_password = st.text_input("Password", type="password")
            if st.button("Login"):
                login_user(login_user_name, login_password)
        else:
            reg_user_name = st.text_input("Choose Username")
            reg_password = st.text_input("Choose Password", type="password")
            reg_full_name = st.text_input("Full Name")
            reg_email = st.text_input("Email")
            reg_company = st.text_input("Company")
            if st.button("Register"):
                register_user(reg_user_name, reg_password, reg_full_name, reg_email, reg_company)
    else:
        st.write(f"Logged in as **{st.session_state['username']}**")
        if st.button("Logout"):
            st.session_state['logged_in'] = False
            st.session_state['user_id'] = None
            st.session_state['username'] = ''
            st.session_state['password'] = ''
            st.success("Logged out")


# Main Title and Executive Header
st.title("📊 Marketing Budget Optimization System")
st.markdown("##### *Enterprise-Grade Linear Programming Framework for Strategic Resource Allocation*")

# Introduction & LaTeX Math Model Section
with st.expander("📚 Mathematical Optimization Framework (LaTeX)", expanded=False):
    st.markdown("""
    ### Linear Programming (LP) Formulation
    To maximize our total brand impressions (Reach) across all social platforms while complying with budgetary boundaries and minimum operational limits, we construct a Linear Programming model.
    
    Let:
    * $x_1$ = Budget allocated to **Instagram** (₹)
    * $x_2$ = Budget allocated to **Google Ads** (₹)
    * $x_3$ = Budget allocated to **YouTube Ads** (₹)
    
    Let $c_1, c_2, c_3$ denote the **Reach Efficiency Coefficients** (Impressions per unit spend) for each respective channel.
    """)
    
    st.latex(r"""
    \begin{aligned}
    \textbf{Maximize Reach (Objective Function):} \quad & Z = c_1 x_1 + c_2 x_2 + c_3 x_3 \\
    \textbf{Subject to:} \quad & x_1 + x_2 + x_3 \le B \quad \text{(Total Budget Constraint)} \\
    & L_1 \le x_1 \le U_1 \quad \text{(Instagram Limits)} \\
    & L_2 \le x_2 \le U_2 \quad \text{(Google Ads Limits)} \\
    & L_3 \le x_3 \le U_3 \quad \text{(YouTube Ads Limits)}
    \end{aligned}
    """)
    
    st.markdown("""
    #### Dual Problem Reformulation for Minimization (SciPy Engine)
    Since `scipy.optimize.linprog` is built exclusively as a minimization solver, we translate our maximization objective function by negating all objective coefficients:
    """)
    
    st.latex(r"""
    \textbf{Minimize:} \quad -Z = -c_1 x_1 - c_2 x_2 - c_3 x_3
    """)
    st.markdown("Upon completion, the solver returns $-Z$, which we negate back ($Z = -(-Z)$) to yield the optimal maximum reach.")

# ==========================================
# SIDEBAR CONTROLS & INPUT MATRIX
# ==========================================
st.sidebar.header("⚙️ Global Configurations")

# Total Budget Input
total_budget = st.sidebar.number_input(
    "Total Available Budget (₹)",
    min_value=1000.0,
    max_value=1000000.0,
    value=10000.0,
    step=500.0,
    help="Define the total maximum monetary cap for reallocation."
)

st.sidebar.markdown("---")
st.sidebar.header("📱 Platform Configurations")

# Instagram parameters
st.sidebar.subheader("📸 Instagram")
ig_c = st.sidebar.number_input("Reach Efficiency (Views/₹) - IG", min_value=1.0, value=50.0, key="ig_c_val")
ig_l = st.sidebar.number_input("Lower Bound Min Spend (₹) - IG", min_value=0.0, value=1000.0, key="ig_l_val")
ig_u = st.sidebar.number_input("Upper Bound Max Spend (₹) - IG", min_value=0.0, value=5000.0, key="ig_u_val")

st.sidebar.markdown("---")

# Google Ads parameters
st.sidebar.subheader("🔍 Google Ads")
go_c = st.sidebar.number_input("Reach Efficiency (Views/₹) - Google", min_value=1.0, value=40.0, key="go_c_val")
go_l = st.sidebar.number_input("Lower Bound Min Spend (₹) - Google", min_value=0.0, value=1500.0, key="go_l_val")
go_u = st.sidebar.number_input("Upper Bound Max Spend (₹) - Google", min_value=0.0, value=4000.0, key="go_u_val")

st.sidebar.markdown("---")

# YouTube Ads parameters
st.sidebar.subheader("🎥 YouTube Ads")
yt_c = st.sidebar.number_input("Reach Efficiency (Views/₹) - YouTube Ads", min_value=1.0, value=60.0, key="yt_c_val")
yt_l = st.sidebar.number_input("Lower Bound Min Spend (₹) - YouTube Ads", min_value=0.0, value=1000.0, key="yt_l_val")
yt_u = st.sidebar.number_input("Upper Bound Max Spend (₹) - YouTube Ads", min_value=0.0, value=6000.0, key="yt_u_val")


# ==========================================
# SOLVER ENGINE & EXCEPTION HANDLING
# ==========================================

# 1. Pre-solver Validation Checks
total_min_required = ig_l + go_l + yt_l
bounds_invalid = (ig_l > ig_u) or (go_l > go_u) or (yt_l > yt_u)

if bounds_invalid:
    st.error("### ⚠️ Invalid Platform Bound Constraints")
    st.markdown(
        f"""
        One or more of your platforms has a **Minimum Operational Spend (Lower Bound)** that exceeds its **Maximum Safety Cap (Upper Bound)**. 
        
        **Please verify your inputs:**
        * **Instagram:** Min: ₹{ig_l:,.2f} | Max: ₹{ig_u:,.2f} {"❌ (Invalid)" if ig_l > ig_u else "✅"}
        * **Google Ads:** Min: ₹{go_l:,.2f} | Max: ₹{go_u:,.2f} {"❌ (Invalid)" if go_l > go_u else "✅"}
        * **YouTube Ads:** Min: ₹{yt_l:,.2f} | Max: ₹{yt_u:,.2f} {"❌ (Invalid)" if yt_l > yt_u else "✅"}
        """
    )
elif total_min_required > total_budget:
    st.error("### ⚠️ Budget Deficit - Infeasible Constraints")
    st.markdown(
        f"""
        The sum of the minimum required operational spend for all platforms is greater than your total available budget!
        
        * **Sum of Platform Lower Bounds ($L_1 + L_2 + L_3$):** **₹{total_min_required:,.2f}**
        * **Your Total Allocated Budget ($B$):** **₹{total_budget:,.2f}**
        * **Required Deficit:** **₹{total_min_required - total_budget:,.2f}**
        
        *To solve this issue, please increase your **Total Available Budget** in the sidebar, or decrease the **Lower Bound Spend** limits of your platforms.*
        """
    )
else:
    # Build LP matrices
    # Objective: Maximize Reach. Correct negation logic for minimization solver:
    # linprog solves: min c^T x  ===>  to max c^T x, we solve: min -c^T x
    c_coefficients = np.array([ig_c, go_c, yt_c])
    c_negated = -c_coefficients
    
    # Inequality constraints (A_ub * x <= b_ub)
    # x1 + x2 + x3 <= Total Budget
    A_ub = np.array([[1.0, 1.0, 1.0]])
    b_ub = np.array([total_budget])
    
    # Platform Bounds [(L1, U1), (L2, U2), (L3, U3)]
    bounds = [
        (ig_l, ig_u),
        (go_l, go_u),
        (yt_l, yt_u)
    ]
    
    try:
        # Solve using HiGHS method
        result = linprog(
            c=c_negated,
            A_ub=A_ub,
            b_ub=b_ub,
            bounds=bounds,
            method='highs'
        )
        
        if result.success:
            x_optimal = result.x
            # Guard against potential microscopic solver noise resulting in negative tiny floats
            x_optimal = np.clip(x_optimal, a_min=[ig_l, go_l, yt_l], a_max=[ig_u, go_u, yt_u])
            
            calculated_reach = -result.fun
            total_allocated_capital = np.sum(x_optimal)
            
            # Localized impressions calculations
            ig_reach = x_optimal[0] * ig_c
            go_reach = x_optimal[1] * go_c
            yt_reach = x_optimal[2] * yt_c
            
            # ==========================================
            # DASHBOARD LAYOUT & ANALYTICS
            # ==========================================
            
            # 1. KPI Cards Row
            kpi_col1, kpi_col2, kpi_col3 = st.columns(3)
            
            with kpi_col1:
                st.metric(
                    label="🎯 Maximum Calculated Reach", 
                    value=f"{int(round(calculated_reach)):,} Impressions",
                    help="The maximum possible reach (impressions) calculated based on optimal spend allocation."
                )
                
            with kpi_col2:
                st.metric(
                    label="💳 Total Allocated Capital", 
                    value=f"₹{total_allocated_capital:,.2f}",
                    delta=f"{((total_allocated_capital / total_budget) * 100):.1f}% Budget Utilized",
                    delta_color="normal",
                    help="Total currency value allocated out of the available budget limit."
                )
                
            with kpi_col3:
                slack = total_budget - total_allocated_capital
                st.metric(
                    label="💰 Unallocated Capital (Slack)",
                    value=f"₹{slack:,.2f}",
                    delta=f"{((slack / total_budget) * 100):.1f}% Leftover",
                    delta_color="inverse",
                    help="Leftover unallocated capital due to maximum budget constraints."
                )
            
            # 2. DataFrame and Plots Section
            tab1, tab2 = st.tabs(["📊 Optimization Analysis", "📋 Technical Model Execution Logs"])
            
            with tab1:
                # Construct clean DataFrame
                platform_names = ["Instagram", "Google Ads", "YouTube Ads"]
                spend_splits = [x_optimal[0], x_optimal[1], x_optimal[2]]
                efficiencies = [ig_c, go_c, yt_c]
                reaches = [ig_reach, go_reach, yt_reach]
                percentages = [(s / total_allocated_capital) * 100 if total_allocated_capital > 0 else 0 for s in spend_splits]
                
                df_profile = pd.DataFrame({
                    "Platform": platform_names,
                    "Optimal Budget Allocation (₹)": spend_splits,
                    "Budget Share (%)": percentages,
                    "Reach Efficiency (Views/₹)": efficiencies,
                    "Expected Impressions (Views)": reaches
                })
                
                # Render Clean DataFrame
                st.markdown("### 📈 Optimal Allocation Data Profile Table")
                try:
                    st.dataframe(
                        df_profile.style.format({
                            "Optimal Budget Allocation (₹)": "₹{:,.2f}",
                            "Budget Share (%)": "{:.2f}%",
                            "Reach Efficiency (Views/₹)": "{:,.1f}",
                            "Expected Impressions (Views)": "{:,.0f}"
                        }),
                        use_container_width=True,
                        hide_index=True
                    )
                except Exception as chart_err:
                    st.warning(f"⚠️ Table rendering issue: {chart_err}")
                    st.write(df_profile)
                
                # Graphical Visualizations in Columns
                chart_col1, chart_col2 = st.columns(2)
                
                with chart_col1:
                    st.markdown("### 📊 Optimal Spend Allocation per Platform")
                    try:
                        # Plotly Corporate Styled Bar Chart
                        fig_bar = go.Figure(data=[
                            go.Bar(
                                x=platform_names,
                                y=spend_splits,
                                text=[f"₹{s:,.0f}" for s in spend_splits],
                                textposition='auto',
                                marker=dict(
                                    color=['#6c5ce7', '#0984e3', '#00cec9'],
                                    line=dict(color='rgba(0, 0, 0, 0.1)', width=1)
                                )
                            )
                        ])
                        fig_bar.update_layout(
                            yaxis_title="Budget Split (₹)",
                            xaxis_title="Platform",
                            template="plotly_white",
                            margin=dict(l=20, r=20, t=20, b=20),
                            height=400,
                            hovermode="x unified"
                        )
                        st.plotly_chart(fig_bar, use_container_width=True)
                    except Exception as chart_err:
                        st.warning(f"⚠️ Bar chart rendering failed: {chart_err}")
                        st.write("**Optimal Spend Allocation:**")
                        for name, spend in zip(platform_names, spend_splits):
                            st.write(f"- {name}: ₹{spend:,.2f}")
                    
                with chart_col2:
                    st.markdown("### 🍕 Percentage Budget Distribution")
                    try:
                        # Plotly Corporate Styled Pie Chart
                        fig_pie = go.Figure(data=[
                            go.Pie(
                                labels=platform_names,
                                values=spend_splits,
                                hole=0.4,
                                marker=dict(
                                    colors=['#6c5ce7', '#0984e3', '#00cec9']
                                ),
                                textinfo='percent+label',
                                insidetextorientation='radial'
                            )
                        ])
                        fig_pie.update_layout(
                            template="plotly_white",
                            margin=dict(l=20, r=20, t=20, b=20),
                            height=400,
                            showlegend=False
                        )
                        st.plotly_chart(fig_pie, use_container_width=True)
                    except Exception as chart_err:
                        st.warning(f"⚠️ Pie chart rendering failed: {chart_err}")
                        st.write("**Budget Distribution:**")
                        for name, spend, pct in zip(platform_names, spend_splits, percentages):
                            st.write(f"- {name}: ₹{spend:,.2f} ({pct:.1f}%)")
                    
            with tab2:
                # Technical Solver summary logs for reviews
                st.markdown("### ⚙️ Solver Execution Diagnostics")
                st.json({
                    "status_code": int(result.status),
                    "status_message": result.message,
                    "solver_method": "HiGHS",
                    "success": bool(result.success),
                    "number_of_iterations": int(result.nit),
                    "raw_slack_variables": list(result.slack),
                    "raw_optimal_x": list(result.x),
                    "raw_optimal_fun": float(result.fun)
                })
                
                st.info(
                    """
                    **Methodology Note:**
                    The optimization utilizes the SciPy simplex/interior-point wrapper method `highs`, 
                    a modern high-performance solver suite designed for large-scale linear and mixed-integer programming problems.
                    """
                )
                
        else:
            # Handle solver failure cases elegantly
            st.error(f"### ⚠️ Optimization Failed")
            st.markdown(
                f"""
                The optimization engine completed execution but could not find a mathematically feasible solution.
                
                **Solver Reason:** `{result.message}`
                
                *This typically occurs if upper bound limits are strictly smaller than the lower bound limits or if constraints create an empty feasibility region. Please check and adjust your platform constraints.*
                """
            )
            
    except Exception as e:
        st.error(f"### ⚠️ Math Optimization Engine Exception")
        st.markdown(
            f"""
            An unexpected error occurred in the Scipy Optimization Engine:
            
            ```python
            Error: {str(e)}
            ```
            
            Please verify that your input fields contain only valid numeric figures and do not present contradictory parameters.
            """
        )