# Constants
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import streamlit as st
from typing import Dict
from utils.llm import summary_generation

ONE_BR_UNITS = 23
TWO_BR_UNITS = 45
SOLAR_PANEL_RATING = 625  # W
BATTERY_CAPACITY = 200  # Ah
BATTERY_VOLTAGE = 96  # V
SYSTEM_LOSSES = 0.20
FEED_IN_TARIFF = 12

# Lighting specifications
LIGHTS_1BR = 5
LIGHTS_2BR = 12
LIGHT_POWER = 6  # Watts per light


def initialize_session_state():
    """Initialize session state variables"""
    defaults = {
        "solar_panels": 20,
        "batteries": 15,
        "panel_price": 13000,
        "battery_price": 39000,
        "grid_price": 28.44,
    }
    for key, value in defaults.items():
        if key not in st.session_state:
            st.session_state[key] = value


def calculate_lighting_consumption(occupancy_1br: float, occupancy_2br: float) -> float:
    """Calculate daily lighting consumption"""
    return (
        (occupancy_1br * ONE_BR_UNITS * LIGHTS_1BR * LIGHT_POWER / 1000)
        + (occupancy_2br * TWO_BR_UNITS * LIGHTS_2BR * LIGHT_POWER / 1000)
    ) * 6  # 6 hours per day


# assuming that 1br average usage is 250wh and for a 2br is 400wh


def calculate_appliance_consumption(
    occupancy_1br: float, occupancy_2br: float
) -> float:
    """Calculate daily appliance consumption by subtracting the lighting usage from the average total consumption for each house type"""
    return (
        occupancy_1br
        * ONE_BR_UNITS
        * (3 - (LIGHTS_1BR * LIGHT_POWER * 6 / 1000))
        # + (500 * 24)  # Fridge
    ) + (
        occupancy_2br
        * TWO_BR_UNITS
        * (4 - (LIGHTS_2BR * LIGHT_POWER * 6 / 1000))
        # + (500 * 24)  # Fridge
    )  # Daily kWh


def total_consumption(
    occupancy_1br: float, occupancy_2br: float, common_area: float
) -> float:
    """Calculate total monthly consumption"""
    lighting = calculate_lighting_consumption(occupancy_1br, occupancy_2br)
    appliances = calculate_appliance_consumption(occupancy_1br, occupancy_2br)
    return (lighting + appliances + common_area) * 30  # Monthly kWh


def solar_production(panels: int) -> float:
    """Monthly solar production with losses"""
    daily_production = (
        panels * SOLAR_PANEL_RATING * 6.5 * (1 - SYSTEM_LOSSES) / 1000
    )  # 6.5 sun hours
    return daily_production * 30  # Monthly kWh


def battery_storage(batteries: int) -> float:
    """Usable battery capacity"""
    return batteries * BATTERY_CAPACITY * BATTERY_VOLTAGE * 0.8 / 1000  # kWh


def financial_analysis(
    consumption: float,
    common_area_consumption: float,
    production: float,
    storage: float,
) -> Dict:
    """Detailed financial calculations"""
    solar_used = min(production, consumption)
    surplus = max(0, production - consumption)
    feed_in_revenue = surplus * FEED_IN_TARIFF / 100  # Convert to Ksh from cents/kWh

    # Account for battery storage
    grid_purchased = max(0, consumption - common_area_consumption - solar_used)
    if storage > 0:
        # Battery can offset some grid purchases
        grid_offset = min(grid_purchased, storage)
        grid_purchased -= grid_offset
    # Money paid to owner if client used this instead o
    monthly_savings = (
        consumption - (common_area_consumption * 30) * st.session_state.grid_price / 100
    )

    total_investment = (
        st.session_state.solar_panels * st.session_state.panel_price
        + st.session_state.batteries * st.session_state.battery_price
    )

    # Avoid division by zero
    if monthly_savings > 0:
        payback_years = total_investment / (monthly_savings * 12)
    else:
        payback_years = float("inf")

    return {
        "consumption": consumption,
        "production": production,
        "solar_contribution": min(100, (solar_used / max(1, consumption)) * 100),
        "grid_dependency": (grid_purchased / max(1, consumption)) * 100,
        "monthly_savings": monthly_savings,
        "payback_period": payback_years,
        "grid_purchased": grid_purchased,
    }


def create_consumption_breakdown(
    occupancy_1br: float, occupancy_2br: float, common_area: float
):
    """Create detailed consumption breakdown"""
    breakdown = {
        "Lighting": calculate_lighting_consumption(occupancy_1br, occupancy_2br) * 30,
        "Appliances": calculate_appliance_consumption(occupancy_1br, occupancy_2br)
        * 30,
        "Common Areas": common_area * 30,
    }
    return pd.DataFrame.from_dict(breakdown, orient="index", columns=["kWh"])


# Streamlit Interface
def main():
    st.set_page_config(page_title="Solar Analysis Suite", page_icon="🌞", layout="wide")
    initialize_session_state()

    # Custom CSS
    st.markdown(
        """
        <style>
        .main .block-container {padding-top: 2rem;}
        h1, h2, h3 {color: #1E88E5;}
        .stExpander {border-radius: 8px; border: 1px solid #1E88E5;}
        .stTabs [data-baseweb="tab-list"] {gap: 10px;}
        .stTabs [data-baseweb="tab"] {
            height: 50px;
            white-space: pre-wrap;
            background-color: #F0F2F6;
            border-radius: 4px 4px 0px 0px;
            gap: 1px;
            padding-top: 10px;
            padding-bottom: 10px;
        }
        .stTabs [aria-selected="true"] {
            background-color: #1E88E5;
            color: white;
        }
        </style>
    """,
        unsafe_allow_html=True,
    )

    # Header with logo
    col1, col2 = st.columns([1, 4])
    with col1:
        st.image("https://img.icons8.com/fluency/96/000000/sun.png", width=100)
    with col2:
        st.title("🌞 Advanced Solar Performance Analyzer")
        st.markdown(
            "Optimize your apartment complex solar installation with data-driven insights"
        )

    # Sidebar for system configuration
    with st.sidebar:
        st.header("System Configuration")

        # Add a nice header image
        st.image("https://img.icons8.com/color/96/000000/solar-panel.png", width=80)

        # Create tabs for different settings
        tab1, tab2 = st.tabs(["Hardware", "Pricing"])

        with tab1:
            st.number_input(
                "Number of Solar Panels",
                1,
                300,
                step=5,
                key="solar_panels",
                help="Each panel rated at 625W",
            )
            st.number_input(
                "Number of Batteries",
                0,
                150,
                step=5,
                key="batteries",
                help="Each battery has 200Ah capacity at 12V",
            )

        with tab2:
            st.number_input(
                "Panel Price (Ksh)",
                1000,
                50000,
                step=500,
                key="panel_price",
                help="Cost per solar panel",
            )
            st.number_input(
                "Battery Price (Ksh)",
                5000,
                100000,
                step=1000,
                key="battery_price",
                help="Cost per battery unit",
            )
            st.number_input(
                "Grid Price (Ksh/kWh)",
                10.0,
                50.0,
                step=0.1,
                key="grid_price",
                help="Current electricity price from the grid",
            )

        st.markdown("---")
        st.markdown(
            """
        📊 **System Totals**
        - **Total Panel Capacity**: {0:.1f} kW
        - **Total Battery Storage**: {1:.1f} kWh
        - **Total Investment**: ksh. {2:,.0f}
          """.format(
                st.session_state.solar_panels * SOLAR_PANEL_RATING / 1000,
                battery_storage(st.session_state.batteries),
                st.session_state.solar_panels * st.session_state.panel_price
                + st.session_state.batteries * st.session_state.battery_price,
            )
        )

    # Main content
    # Create scenarios with varying occupancy levels
    scenarios = {}

    # Common area consumption remains constant
    common_area_consumption = 23.544  # kWh per day

    # Generate scenarios with different occupancy combinations
    occupancy_levels = [0.0, 0.25, 0.50, 0.75, 1.0]

    # Create scenarios for 1BR fixed, varying 2BR
    for br1_level in occupancy_levels:
        for br2_level in occupancy_levels:
            scenario_name = f"1BR: {int(br1_level*100)}%, 2BR: {int(br2_level*100)}%"
            scenarios[scenario_name] = {
                "1br": br1_level,
                "2br": br2_level,
                "common": common_area_consumption,
            }

    # Analysis tabs
    st.markdown("---")
    tab1, tab2, tab3 = st.tabs(
        ["📊 Energy Analysis", "💰 Financial Metrics", "🔍 Detailed Breakdown"]
    )

    # Prepare analysis data for all scenarios
    analysis_data = []
    for name, params in scenarios.items():
        consumption = total_consumption(params["1br"], params["2br"], params["common"])
        production = solar_production(st.session_state.solar_panels)
        storage = battery_storage(st.session_state.batteries)
        financials = financial_analysis(
            consumption, common_area_consumption, production, storage
        )
        analysis_data.append({"Scenario": name, **financials})

    df = pd.DataFrame(analysis_data)

    # Tab 1: Energy Analysis
    with tab1:
        st.header("Energy Flow Analysis")

        # Allow filtering by 1BR occupancy
        one_br_filter = st.selectbox(
            "Filter by 1BR Occupancy",
            ["All"] + [f"{int(level*100)}%" for level in occupancy_levels],
            help="Filter scenarios by 1BR occupancy level",
        )

        # Filter the dataframe based on selection
        filtered_df = df
        if one_br_filter != "All":
            occupancy_value = int(one_br_filter.replace("%", ""))
            filtered_df = df[df["Scenario"].str.contains(f"1BR: {occupancy_value}%")]

        # Chart 1: Energy Balance
        st.subheader("Energy Balance by Scenario")

        energy_fig = plt.figure(figsize=(12, 7))
        ax = energy_fig.add_subplot(111)

        # Create data for stacked bar chart
        chart_data = filtered_df.copy()
        chart_data["grid_energy"] = chart_data["grid_purchased"]
        chart_data["solar_energy"] = (
            chart_data["consumption"] - chart_data["grid_purchased"]
        )

        # Create normalized stacked bar chart
        chart_data = chart_data.set_index("Scenario")
        energy_proportions = (
            chart_data[["solar_energy", "grid_energy"]].div(
                chart_data["consumption"], axis=0
            )
            * 100
        )
        energy_proportions = energy_proportions.reset_index()

        # Reshape for seaborn
        energy_melt = pd.melt(
            energy_proportions,
            id_vars=["Scenario"],
            value_vars=["solar_energy", "grid_energy"],
            var_name="Energy Source",
            value_name="Percentage",
        )

        # Rename for better labels
        energy_melt["Energy Source"] = energy_melt["Energy Source"].replace(
            {"solar_energy": "Solar Generated", "grid_energy": "Grid Purchased"}
        )

        # Plot with seaborn
        sns.set_theme(style="whitegrid")
        sns.barplot(
            data=energy_melt,
            x="Scenario",
            y="Percentage",
            hue="Energy Source",
            palette=["#4CAF50", "#F44336"],
            ax=ax,
        )
        ax.set_ylabel("Energy Contribution (%)")
        ax.set_title("Energy Source Distribution by Occupancy Scenario")
        plt.xticks(rotation=45, ha="right")
        plt.tight_layout()
        st.pyplot(energy_fig)

        # Detailed metrics
        col1, col2, col3 = st.columns(3)
        with col1:
            st.metric(
                "Avg. Solar Contribution",
                f"{filtered_df['solar_contribution'].mean():.1f}%",
                (
                    f"{filtered_df['solar_contribution'].mean() - 50:.1f}%"
                    if filtered_df["solar_contribution"].mean() > 50
                    else f"{filtered_df['solar_contribution'].mean() - 50:.1f}%"
                ),
            )
        with col2:
            st.metric(
                "Avg. Grid Dependency",
                f"{filtered_df['grid_dependency'].mean():.1f}%",
                (
                    f"{50 - filtered_df['grid_dependency'].mean():.1f}%"
                    if filtered_df["grid_dependency"].mean() < 50
                    else f"{50 - filtered_df['grid_dependency'].mean():.1f}%"
                ),
            )
        with col3:
            st.metric(
                "Production/Consumption Ratio",
                f"{(filtered_df['production'].mean() / filtered_df['consumption'].mean() * 100):.1f}%",
            )

        with st.expander("🔍 Energy Flow Interpretation"):
            st.markdown(
                """
                **Understanding the Chart:**
                - **Solar Contribution**: Percentage of total energy needs met directly by solar production
                - **Grid Dependency**: Remaining energy required from the grid
                - The ideal scenario shows high solar contribution with minimal grid dependency
                
                **Key Factors Affecting Energy Balance:**
                1. **Occupancy Levels**: Higher occupancy means higher consumption, which may exceed solar capacity
                2. **Solar System Size**: More panels increase production and reduce grid dependency
                3. **Battery Storage**: Helps utilize excess daytime production for nighttime use
                """
            )

    # Tab 2: Financial Metrics
    with tab2:
        st.header("Financial Performance Analysis")

        # Allow filtering by 2BR occupancy
        two_br_filter = st.selectbox(
            "Filter by 2BR Occupancy",
            ["All"] + [f"{int(level*100)}%" for level in occupancy_levels],
            help="Filter scenarios by 2BR occupancy level",
        )

        # Filter the dataframe based on selection
        filtered_fin_df = df
        if two_br_filter != "All":
            occupancy_value = int(two_br_filter.replace("%", ""))
            filtered_fin_df = df[
                df["Scenario"].str.contains(f"2BR: {occupancy_value}%")
            ]

        # Monthly Savings Chart
        st.subheader("Monthly Cost Savings")

        # Fix large values
        filtered_fin_df["monthly_savings_fixed"] = filtered_fin_df[
            "monthly_savings"
        ].clip(0, 100000)

        fig1, ax1 = plt.subplots(figsize=(12, 6))
        sns.barplot(
            data=filtered_fin_df,
            x="Scenario",
            y="monthly_savings_fixed",
            palette="viridis",
            ax=ax1,
        )
        ax1.set_title("Monthly Cost Savings by Scenario")
        ax1.set_ylabel("Ksh")
        plt.xticks(rotation=45, ha="right")
        plt.tight_layout()
        st.pyplot(fig1)

        # Payback Period Chart
        st.subheader("System Payback Period")

        # Fix large values
        filtered_fin_df["payback_period_fixed"] = filtered_fin_df[
            "payback_period"
        ].clip(0, 30)

        fig2, ax2 = plt.subplots(figsize=(12, 6))
        sns.barplot(
            data=filtered_fin_df,
            x="Scenario",
            y="payback_period_fixed",
            palette="rocket_r",
            ax=ax2,
        )
        ax2.set_title("Investment Payback Period by Scenario")
        ax2.set_ylabel("Years")
        plt.xticks(rotation=45, ha="right")
        plt.tight_layout()
        st.pyplot(fig2)

        # Financial summary metrics
        col1, col2, col3 = st.columns(3)
        with col1:
            avg_savings = filtered_fin_df["monthly_savings"].mean()
            st.metric(
                "Avg. Monthly Savings",
                f"{avg_savings:,.0f} Ksh",
                (
                    f"{avg_savings - df['monthly_savings'].mean():,.0f} Ksh"
                    if avg_savings > df["monthly_savings"].mean()
                    else f"{avg_savings - df['monthly_savings'].mean():,.0f} Ksh"
                ),
            )
        with col2:
            min_payback = filtered_fin_df["payback_period"].min()
            st.metric(
                "Best Payback Period",
                f"{min_payback:.1f} years",
                help="Shortest time to recover investment",
            )
        with col3:
            total_investment = (
                st.session_state.solar_panels * st.session_state.panel_price
                + st.session_state.batteries * st.session_state.battery_price
            )
            annual_roi = (
                (avg_savings * 12 / total_investment) * 100
                if total_investment > 0
                else 0
            )
            st.metric(
                "Annual ROI", f"{annual_roi:.1f}%", help="Annual Return on Investment"
            )

        with st.expander("💵 Financial Analysis Details"):
            st.markdown(
                f"""
                **Investment Details:**
                - Total Solar Panel Investment: {st.session_state.solar_panels:,} panels × {st.session_state.panel_price:,} Ksh = {st.session_state.solar_panels * st.session_state.panel_price:,} Ksh
                - Total Battery Investment: {st.session_state.batteries:,} batteries × {st.session_state.battery_price:,} Ksh = {st.session_state.batteries * st.session_state.battery_price:,} Ksh
                - Total System Cost: {total_investment:,} Ksh
                
                **Savings Calculation:**
                - Grid Price: {st.session_state.grid_price} Ksh/kWh
                - Monthly Savings =(Total Consumption - Common Area) × Grid Price
                - Payback Period = Total Investment / Annual Savings
                
                **Filtered Scenario Data:**
                """
            )
            st.dataframe(
                filtered_fin_df[
                    [
                        "Scenario",
                        "consumption",
                        "production",
                        "monthly_savings",
                        "payback_period",
                    ]
                ].sort_values("monthly_savings", ascending=False),
                hide_index=True,
            )
            # Button to trigger analysis
            if st.button("🔍 Analyze Financial Data with LLM"):
                with st.spinner("Generating insights with AI..."):
                    analysis = summary_generation(filtered_fin_df)
                    st.success("Analysis Complete!")
                    st.write(analysis)  # Display the results

    # Tab 3: Detailed Breakdown
    with tab3:
        st.header("Consumption Breakdown Analysis")

        # Select specific scenario for detailed analysis
        scenario_select = st.selectbox(
            "Select Specific Scenario", list(scenarios.keys())
        )
        selected_params = scenarios[scenario_select]

        # Create consumption breakdown
        breakdown_df = create_consumption_breakdown(
            selected_params["1br"], selected_params["2br"], selected_params["common"]
        )

        total_kwh = breakdown_df["kWh"].sum()

        # Add percentage column
        breakdown_df["Percentage"] = (breakdown_df["kWh"] / total_kwh * 100).round(1)

        col1, col2 = st.columns([2, 3])

        with col1:
            st.subheader("Energy Composition")

            # Create a more attractive pie chart
            fig3 = plt.figure(figsize=(8, 8))
            ax3 = fig3.add_subplot(111)

            colors = ["#FF9800", "#2196F3", "#4CAF50"]
            explode = (0.1, 0, 0)

            wedges, texts, autotexts = ax3.pie(
                breakdown_df["kWh"],
                labels=breakdown_df.index,
                autopct="%1.1f%%",
                explode=explode,
                colors=colors,
                shadow=True,
                startangle=90,
                textprops={"fontsize": 12},
            )

            # Equal aspect ratio ensures that pie is drawn as a circle
            ax3.axis("equal")
            plt.tight_layout()
            st.pyplot(fig3)

            # Show total consumption
            st.metric(
                "Total Monthly Consumption",
                f"{total_kwh:.1f} kWh",
                help="Sum of all consumption components",
            )

        with col2:
            st.subheader("Detailed Component Analysis")

            # Show breakdown as a horizontal bar chart
            fig4 = plt.figure(figsize=(10, 5))
            ax4 = fig4.add_subplot(111)

            # Sort by consumption
            sorted_df = breakdown_df.sort_values("kWh", ascending=True)

            # Create horizontal bar chart
            bars = sns.barplot(
                y=sorted_df.index, x="kWh", data=sorted_df, palette=colors[::-1], ax=ax4
            )

            # Add data labels
            for i, v in enumerate(sorted_df["kWh"]):
                ax4.text(
                    v + 5,
                    i,
                    f"{v:.1f} kWh ({sorted_df['Percentage'].iloc[i]}%)",
                    va="center",
                )

            ax4.set_title(f"Energy Consumption Breakdown - {scenario_select}")
            ax4.set_xlabel("Monthly Consumption (kWh)")
            ax4.set_ylabel("")
            plt.tight_layout()
            st.pyplot(fig4)

            # Add scenario details
            st.markdown(
                f"""
            **Scenario Details:**
            - 1BR Units Occupancy: {selected_params['1br']*100:.0f}% ({selected_params['1br']*ONE_BR_UNITS:.0f} units)
            - 2BR Units Occupancy: {selected_params['2br']*100:.0f}% ({selected_params['2br']*TWO_BR_UNITS:.0f} units)
            - Common Areas Consumption: {selected_params['common']*30:.1f} kWh/month
            """
            )

            # Insight box
            st.info(
                f"""
            **Key Insights for {scenario_select}:**
            - Lighting contributes {breakdown_df.loc['Lighting', 'Percentage']:.1f}% of total consumption
            - Common areas account for {breakdown_df.loc['Common Areas', 'Percentage']:.1f}% of the total
            - {'2BR units dominate consumption at ' + str(selected_params['2br']*100) + '% occupancy' if selected_params['2br'] > selected_params['1br'] else '1BR units are the primary consumers at ' + str(selected_params['1br']*100) + '% occupancy'}
            - Total potential solar offset: {min(solar_production(st.session_state.solar_panels)/total_kwh*100, 100):.1f}%
            """
            )

    # Footer
    st.markdown("---")
    st.markdown(
        """
        <div style="text-align: center; color: #666;">
        <p>Solar Analysis Suite v1.0 | Developed with ❤️ for sustainable energy solutions</p>
        </div>
        """,
        unsafe_allow_html=True,
    )


if __name__ == "__main__":
    main()