import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
from sys_of_eqn_solver import SystemSolver
import random

# Page configuration
st.set_page_config(
    page_title="Systems of Equations - Algebra II",
    page_icon="📊",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize session state
if 'solver' not in st.session_state:
    st.session_state.solver = SystemSolver()
if 'show_steps' not in st.session_state:
    st.session_state.show_steps = True

# Header
st.title("📊 Systems of Linear Equations - Algebra II Tutorial")
st.markdown("*Learn to solve systems of equations using multiple methods with real-time visualization*")

# Sidebar configuration
with st.sidebar:
    st.header("⚙️ Configuration")
    
    # System size selection
    system_size = st.selectbox(
        "Select System Size:",
        ["2x2 System", "3x3 System"],
        help="Choose between 2 equations with 2 variables or 3 equations with 3 variables"
    )
    
    # Method selection
    if system_size == "2x2 System":
        available_methods = ["All Methods", "Graphical", "Substitution", "Elimination", "Matrix"]
    else:
        available_methods = ["All Methods", "Matrix", "Elimination"]
    
    selected_method = st.selectbox("Solution Method:", available_methods)
    
    # Display options
    st.subheader("Display Options")
    show_steps = st.checkbox("Show detailed steps", value=st.session_state.show_steps)
    show_verification = st.checkbox("Show solution verification", value=True)
    
    if system_size == "2x2 System":
        show_graph = st.checkbox("Show graphical solution", value=True)
        graph_range = st.slider("Graph range", min_value=5, max_value=20, value=10)
    
    # Quick examples
    st.subheader("📚 Quick Examples")
    if st.button("🎯 Unique Solution"):
        if system_size == "2x2 System":
            st.session_state.example_coeffs = [[2, 1], [1, -1]]
            st.session_state.example_constants = [7, 1]
        else:
            st.session_state.example_coeffs = [[1, 2, -1], [2, 1, 1], [1, -1, 2]]
            st.session_state.example_constants = [3, 7, 4]
    
    if st.button("🔄 Infinite Solutions"):
        if system_size == "2x2 System":
            st.session_state.example_coeffs = [[2, 4], [1, 2]]
            st.session_state.example_constants = [6, 3]
        else:
            st.session_state.example_coeffs = [[1, 2, 3], [2, 4, 6], [1, 2, 3]]
            st.session_state.example_constants = [6, 12, 6]
    
    if st.button("❌ No Solution"):
        if system_size == "2x2 System":
            st.session_state.example_coeffs = [[2, 4], [1, 2]]
            st.session_state.example_constants = [6, 4]
        else:
            st.session_state.example_coeffs = [[1, 2, 3], [2, 4, 6], [1, 2, 3]]
            st.session_state.example_constants = [6, 12, 7]
    
    if st.button("🎲 Random System"):
        if system_size == "2x2 System":
            st.session_state.example_coeffs = [[random.randint(1, 5), random.randint(1, 5)], 
                                             [random.randint(1, 5), random.randint(1, 5)]]
            st.session_state.example_constants = [random.randint(1, 15), random.randint(1, 15)]
        else:
            st.session_state.example_coeffs = [[random.randint(1, 4), random.randint(1, 4), random.randint(1, 4)],
                                             [random.randint(1, 4), random.randint(1, 4), random.randint(1, 4)],
                                             [random.randint(1, 4), random.randint(1, 4), random.randint(1, 4)]]
            st.session_state.example_constants = [random.randint(1, 12), random.randint(1, 12), random.randint(1, 12)]

# Main content area
if system_size == "2x2 System":
    st.header("2×2 System of Linear Equations")
    
    # Input section
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("📝 Enter Your System")
        
        # Use example values if they exist
        if 'example_coeffs' in st.session_state and len(st.session_state.example_coeffs) == 2:
            default_coeffs = st.session_state.example_coeffs
            default_constants = st.session_state.example_constants
        else:
            default_coeffs = [[2, 1], [1, -1]]
            default_constants = [7, 1]
        
        # Equation 1
        st.markdown("**Equation 1:** `a₁x + b₁y = c₁`")
        eq1_col1, eq1_col2, eq1_col3 = st.columns(3)
        with eq1_col1:
            a1 = st.number_input("a₁", value=float(default_coeffs[0][0]), step=0.1, key="a1")
        with eq1_col2:
            b1 = st.number_input("b₁", value=float(default_coeffs[0][1]), step=0.1, key="b1")
        with eq1_col3:
            c1 = st.number_input("c₁", value=float(default_constants[0]), step=0.1, key="c1")
        
        # Equation 2
        st.markdown("**Equation 2:** `a₂x + b₂y = c₂`")
        eq2_col1, eq2_col2, eq2_col3 = st.columns(3)
        with eq2_col1:
            a2 = st.number_input("a₂", value=float(default_coeffs[1][0]), step=0.1, key="a2")
        with eq2_col2:
            b2 = st.number_input("b₂", value=float(default_coeffs[1][1]), step=0.1, key="b2")
        with eq2_col3:
            c2 = st.number_input("c₂", value=float(default_constants[1]), step=0.1, key="c2")
        
        # Display the system
        st.markdown("### Your System:")
        if a1 >= 0 and b1 >= 0:
            eq1_str = f"{a1}x + {b1}y = {c1}"
        elif a1 >= 0 and b1 < 0:
            eq1_str = f"{a1}x - {abs(b1)}y = {c1}"
        elif a1 < 0 and b1 >= 0:
            eq1_str = f"-{abs(a1)}x + {b1}y = {c1}"
        else:
            eq1_str = f"-{abs(a1)}x - {abs(b1)}y = {c1}"
        
        if a2 >= 0 and b2 >= 0:
            eq2_str = f"{a2}x + {b2}y = {c2}"
        elif a2 >= 0 and b2 < 0:
            eq2_str = f"{a2}x - {abs(b2)}y = {c2}"
        elif a2 < 0 and b2 >= 0:
            eq2_str = f"-{abs(a2)}x + {b2}y = {c2}"
        else:
            eq2_str = f"-{abs(a2)}x - {abs(b2)}y = {c2}"
        
        st.latex(eq1_str)
        st.latex(eq2_str)
    
    with col2:
        # Solve the system
        coefficients = [[a1, b1], [a2, b2]]
        constants = [c1, c2]
        
        method_map = {
            "All Methods": "all",
            "Graphical": "graphical",
            "Substitution": "substitution", 
            "Elimination": "elimination",
            "Matrix": "matrix"
        }
        
        result = st.session_state.solver.solve_2x2_system(
            coefficients, constants, method_map[selected_method]
        )
        
        # System type indicator
        system_type = result['system_type']
        if system_type == "unique_solution":
            st.success("✅ **Unique Solution** (Consistent Independent)")
        elif system_type == "infinite_solutions":
            st.info("🔄 **Infinite Solutions** (Consistent Dependent)")
        else:
            st.error("❌ **No Solution** (Inconsistent)")
        
        # Display solutions based on selected method
        if selected_method != "All Methods":
            method_key = f"{method_map[selected_method]}_solution"
            if method_key in result:
                solution_data = result[method_key]
                
                st.subheader(f"🔍 {selected_method} Method")
                
                if show_steps and 'steps' in solution_data:
                    with st.expander("Show Detailed Steps", expanded=True):
                        for i, step in enumerate(solution_data['steps']):
                            st.write(f"{i+1}. {step}")
                
                if 'solution' in solution_data and solution_data['solution'] is not None:
                    if isinstance(solution_data['solution'], dict):
                        st.markdown("### 🎯 Solution:")
                        col1, col2 = st.columns(2)
                        with col1:
                            st.metric("x =", f"{solution_data['solution']['x']:.4f}")
                        with col2:
                            st.metric("y =", f"{solution_data['solution']['y']:.4f}")
    
    # Graphical visualization
    if show_graph and ('graphical_solution' in result or selected_method == "All Methods"):
        st.header("📈 Graphical Solution")
        
        fig = go.Figure()
        
        # Generate x values for plotting
        x_vals = np.linspace(-graph_range, graph_range, 400)
        
        # Plot first line
        if b1 != 0:
            y1_vals = (c1 - a1 * x_vals) / b1
            fig.add_trace(go.Scatter(
                x=x_vals, y=y1_vals,
                mode='lines',
                name=eq1_str,
                line=dict(color='blue', width=3)
            ))
        else:
            # Vertical line
            x_vert = c1 / a1 if a1 != 0 else 0
            fig.add_vline(x=x_vert, line_dash="solid", line_color="blue", 
                         annotation_text=f"x = {x_vert:.2f}")
        
        # Plot second line
        if b2 != 0:
            y2_vals = (c2 - a2 * x_vals) / b2
            fig.add_trace(go.Scatter(
                x=x_vals, y=y2_vals,
                mode='lines',
                name=eq2_str,
                line=dict(color='red', width=3)
            ))
        else:
            # Vertical line
            x_vert = c2 / a2 if a2 != 0 else 0
            fig.add_vline(x=x_vert, line_dash="solid", line_color="red",
                         annotation_text=f"x = {x_vert:.2f}")
        
        # Plot intersection point
        if 'graphical_solution' in result and result['graphical_solution']['intersection']:
            intersection = result['graphical_solution']['intersection']
            fig.add_trace(go.Scatter(
                x=[intersection['x']], y=[intersection['y']],
                mode='markers',
                name=f"Solution: ({intersection['x']:.3f}, {intersection['y']:.3f})",
                marker=dict(color='green', size=12, symbol='circle')
            ))
        
        fig.update_layout(
            title="Graphical Solution of System",
            xaxis_title="x",
            yaxis_title="y",
            xaxis=dict(range=[-graph_range, graph_range], gridcolor='lightgray'),
            yaxis=dict(range=[-graph_range, graph_range], gridcolor='lightgray'),
            plot_bgcolor='white',
            showlegend=True
        )
        
        st.plotly_chart(fig, use_container_width=True)

else:  # 3x3 System
    st.header("3×3 System of Linear Equations")
    
    # Input section for 3x3 system
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.subheader("📝 Enter Your 3×3 System")
        
        # Use example values if they exist
        if 'example_coeffs' in st.session_state and len(st.session_state.example_coeffs) == 3:
            default_coeffs = st.session_state.example_coeffs
            default_constants = st.session_state.example_constants
        else:
            default_coeffs = [[1, 2, -1], [2, 1, 1], [1, -1, 2]]
            default_constants = [3, 7, 4]
        
        # Create input grid for 3x3 system
        st.markdown("**System of equations:** `Ax = b`")
        
        equations = []
        coefficients = []
        constants = []
        
        for i in range(3):
            st.markdown(f"**Equation {i+1}:** `a{i+1}₁x + a{i+1}₂y + a{i+1}₃z = b{i+1}`")
            cols = st.columns(4)
            
            row_coeffs = []
            with cols[0]:
                val = st.number_input(f"a{i+1}₁", value=float(default_coeffs[i][0]), step=0.1, key=f"a{i+1}1")
                row_coeffs.append(val)
            with cols[1]:
                val = st.number_input(f"a{i+1}₂", value=float(default_coeffs[i][1]), step=0.1, key=f"a{i+1}2")
                row_coeffs.append(val)
            with cols[2]:
                val = st.number_input(f"a{i+1}₃", value=float(default_coeffs[i][2]), step=0.1, key=f"a{i+1}3")
                row_coeffs.append(val)
            with cols[3]:
                val = st.number_input(f"b{i+1}", value=float(default_constants[i]), step=0.1, key=f"b{i+1}")
                constants.append(val)
            
            coefficients.append(row_coeffs)
            
            # Format equation string
            eq_parts = []
            for j, coeff in enumerate(row_coeffs):
                var = ['x', 'y', 'z'][j]
                if j == 0:
                    eq_parts.append(f"{coeff}{var}")
                else:
                    if coeff >= 0:
                        eq_parts.append(f" + {coeff}{var}")
                    else:
                        eq_parts.append(f" - {abs(coeff)}{var}")
            equations.append(''.join(eq_parts) + f" = {constants[i]}")
        
        st.markdown("### Your System:")
        for eq in equations:
            st.latex(eq)
    
    with col2:
        # Solve 3x3 system
        method_map = {"All Methods": "all", "Matrix": "matrix", "Elimination": "elimination"}
        
        result = st.session_state.solver.solve_3x3_system(
            coefficients, constants, method_map[selected_method]
        )
        
        # System type indicator
        system_type = result['system_type']
        if system_type == "unique_solution":
            st.success("✅ **Unique Solution** (Consistent Independent)")
        elif system_type == "infinite_solutions":
            st.info("🔄 **Infinite Solutions** (Consistent Dependent)")
        else:
            st.error("❌ **No Solution** (Inconsistent)")
        
        # Display solutions
        if selected_method != "All Methods":
            method_key = f"{method_map[selected_method]}_solution"
            if method_key in result:
                solution_data = result[method_key]
                
                st.subheader(f"🔍 {selected_method} Method")
                
                if show_steps and 'steps' in solution_data:
                    with st.expander("Show Detailed Steps", expanded=True):
                        for i, step in enumerate(solution_data['steps']):
                            st.write(f"{i+1}. {step}")
                
                if 'solution' in solution_data and solution_data['solution'] is not None:
                    if isinstance(solution_data['solution'], dict):
                        st.markdown("### 🎯 Solution:")
                        col1, col2, col3 = st.columns(3)
                        with col1:
                            st.metric("x =", f"{solution_data['solution']['x']:.4f}")
                        with col2:
                            st.metric("y =", f"{solution_data['solution']['y']:.4f}")
                        with col3:
                            st.metric("z =", f"{solution_data['solution']['z']:.4f}")

# Show all methods if selected
if selected_method == "All Methods":
    st.header("🔄 All Solution Methods")
    
    methods_to_show = []
    if system_size == "2x2 System":
        methods_to_show = ['matrix_solution', 'elimination_solution', 'substitution_solution', 'graphical_solution']
    else:
        methods_to_show = ['matrix_solution', 'elimination_solution']
    
    tabs = st.tabs([method.replace('_solution', '').title() for method in methods_to_show])
    
    for i, method_key in enumerate(methods_to_show):
        with tabs[i]:
            if method_key in result:
                solution_data = result[method_key]
                
                if 'error' in solution_data:
                    st.error(f"Error: {solution_data['error']}")
                else:
                    if show_steps and 'steps' in solution_data:
                        st.subheader("📋 Steps:")
                        for j, step in enumerate(solution_data['steps']):
                            st.write(f"{j+1}. {step}")
                    
                    if 'solution' in solution_data and solution_data['solution'] is not None:
                        if isinstance(solution_data['solution'], dict):
                            st.subheader("🎯 Solution:")
                            if system_size == "2x2 System":
                                col1, col2 = st.columns(2)
                                with col1:
                                    st.metric("x =", f"{solution_data['solution']['x']:.4f}")
                                with col2:
                                    st.metric("y =", f"{solution_data['solution']['y']:.4f}")
                            else:
                                col1, col2, col3 = st.columns(3)
                                with col1:
                                    st.metric("x =", f"{solution_data['solution']['x']:.4f}")
                                with col2:
                                    st.metric("y =", f"{solution_data['solution']['y']:.4f}")
                                with col3:
                                    st.metric("z =", f"{solution_data['solution']['z']:.4f}")

# Solution verification
if show_verification and system_size == "2x2 System":
    if 'matrix_solution' in result and result['matrix_solution']['solution']:
        st.header("✅ Solution Verification")
        
        solution = result['matrix_solution']['solution']
        x_val, y_val = solution['x'], solution['y']
        
        # Check each equation
        check1 = a1 * x_val + b1 * y_val
        check2 = a2 * x_val + b2 * y_val
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Equation 1 Check:")
            st.write(f"{a1}({x_val:.4f}) + {b1}({y_val:.4f}) = {check1:.4f}")
            if abs(check1 - c1) < 1e-6:
                st.success(f"✅ {check1:.4f} = {c1} ✓")
            else:
                st.error(f"❌ {check1:.4f} ≠ {c1}")
        
        with col2:
            st.subheader("Equation 2 Check:")
            st.write(f"{a2}({x_val:.4f}) + {b2}({y_val:.4f}) = {check2:.4f}")
            if abs(check2 - c2) < 1e-6:
                st.success(f"✅ {check2:.4f} = {c2} ✓")
            else:
                st.error(f"❌ {check2:.4f} ≠ {c2}")

# Educational notes
with st.expander("📚 Educational Notes"):
    if system_size == "2x2 System":
        st.markdown("""
        ### 2×2 Systems - Key Concepts:
        
        **Types of Solutions:**
        - **Unique Solution**: Lines intersect at exactly one point (det ≠ 0)
        - **Infinite Solutions**: Lines are identical (same slope, same y-intercept)
        - **No Solution**: Lines are parallel (same slope, different y-intercepts)
        
        **Solution Methods:**
        1. **Graphical**: Plot both lines and find intersection point
        2. **Substitution**: Solve one equation for a variable, substitute into the other
        3. **Elimination**: Add/subtract equations to eliminate a variable
        4. **Matrix/Cramer's Rule**: Use determinants to find solution
        
        **When to Use Each Method:**
        - **Graphical**: Good for visualization and understanding
        - **Substitution**: Best when one variable has coefficient 1 or -1
        - **Elimination**: Good when coefficients are easy to work with
        - **Matrix**: Most systematic, especially for larger systems
        """)
    else:
        st.markdown("""
        ### 3×3 Systems - Key Concepts:
        
        **Types of Solutions:**
        - **Unique Solution**: System has exactly one solution (det ≠ 0)
        - **Infinite Solutions**: rank(A) = rank([A|b]) < 3
        - **No Solution**: rank(A) ≠ rank([A|b])
        
        **Solution Methods:**
        1. **Matrix Method**: Use matrix operations and determinants
        2. **Gaussian Elimination**: Systematic row operations to solve
        
        **Key Skills:**
        - Forward elimination to create upper triangular matrix
        - Back substitution to find solution
        - Recognition of inconsistent and dependent systems
        """)

# Footer
st.markdown("---")
st.markdown("*Built with Streamlit for Algebra II students • Use the sidebar to explore different examples and methods*")