import gradio as gr
import pandas as pd
import numpy as np
from numpy.typing import NDArray
from pyscipopt import Model, quicksum
from datetime import datetime

# Define the quality tiers and names for the plants
PLANTS_TIERS = {
    "radiant": "RADIANT",
    "flourishing": "FLOURISHING",
    "hardy": "HARDY",
    "feeble": "FEEBLE",
    "radiant_rarecolor": "RADIANT+RARE",
    "flourishing_rarecolor": "FLOURISHING+RARE",
    "hardy_rarecolor": "HARDY+RARE",
}

PLANTS_LABLES = {
    "fanged_geranium": "Fanged Geranium",
    "gillyweed": "Gillyweed",
    "rose": "Rose",
    "puffapod": "Puffapod",
    "wild_pansy": "Wild Pansy",
    "nifflers_fancy": "Niffler's Fancy",
    "fanwort": "Fanwort",
    "ladys_mantle": "Lady's Mantle",
    "kelp": "Kelp",
    "mandrake": "Mandrake",
    "chinese_chomping_cabbage": "Chinese Chomping Cabbage",
    "dragons_breath_macroalgae": "Dragon's Breath Macroalgae",
    "peony": "Peony",
    "begonia": "Begonia",
    "mayflower": "Mayflower",
    "hydrangea": "Hydrangea",
    "ludwigia_glandulosa": "Ludwigia Glandulosa",
    "daffodil": "Daffodil",
    "water_hyacinth": "Water Hyacinth",
    "lily_of_the_valley": "Lily of the Valley",
    "mosaic_flower": "Mosaic Flower",
    "sunflower": "Sunflower",
    "mimbulus_mimbletonia": "Mimbulus Mimbletonia",
    "water_lily": "Water Lily",
}

INTERFACE_TEXTS = {
    "cn": {
        "gold_label": "葭碧の金币预算：",
        "strategies_label": "请选择凑单策略：",
        "clear_btn_label": "❌清除",
        "calculate_btn_label": "🛠计算",
        "output_label": "计算结果：",
        "strategy_options": [
            ("最小化售出株数(优先出售高价植物)", "MaximizeStock"),
            ("最大化售出株数(优先出售低价植物)", "MinimizeStock"),
        ],
    },
    "en": {
        "gold_label": "Gabby's Gold Budget:",
        "strategies_label": "Select a strategy:",
        "clear_btn_label": "❌Clear",
        "calculate_btn_label": "🛠Calculate",
        "output_label": "Output:",
        "strategy_options": [
            (
                "Minimize the number of plants sold (prioritize high-priced plants)",
                "MaximizeStock",
            ),
            (
                "Maximize the number of plants sold (prioritize low-priced plants)",
                "MinimizeStock",
            ),
        ],
    },
}


def process_csv(file_path):
    """import and process plants data"""
    df = pd.read_csv(file_path)
    df["species"] = pd.Categorical(df["species"])
    df["tier"] = pd.Categorical(df["tier"])
    # df = df.dropna(subset=["gold"])
    df = df.astype(
        {
            "gold": int,
            "gems": int,
        }
    )
    return df


df = process_csv("plants.csv")

GOLD_PLANTS = set(
    row["species"]
    for _, row in df.iterrows()
    if row["gold"] != 0 and row["tier"] != "feeble"
)
GEMS_PLANTS = set(
    row["species"]
    for _, row in df.iterrows()
    if row["gems"] != 0 and row["tier"] != "feeble"
)


def check_currency(plant, currency):
    if currency == "gold":
        return plant in GOLD_PLANTS
    elif currency == "gems":
        return plant in GEMS_PLANTS


def calculator(currency, budget, strategy, extra_rate, *amount):
    """
    Calculate the optimal solution of plant sales based on the given budget
    and inventory constraints.

    Args:
        *args (tuple): A tuple containing:
            - currency (str): The currency used for purchasing plants ("gold" or "gems").
            - budget (int): Gabby's gold budget.
            - strategy (str): The selected strategy for selling plants ("MaximizeStock" or "MinimizeStock").
            - extra_rate (int): The premium rate for selling plants.
            - amount (list of int): Stock levels of each plant type.

    Returns:
        str: A description of the optimal solution, including which plants to sell,
             the total gold earned, and the remaining inventory.
             Returns an error message if no solution is found.
    """
    # currency: str, budget: int, strategy:str, extra_rate:int = args[0:4]
    # budget: int = args[0]  # 葭碧预算
    # strategy: str = args[1]  # 出售策略
    # extra_rate: int = args[2]  # 高价收购倍率
    stocks: NDArray[np.int_ | np.integer] = np.array(
        [x if x else 0 for x in amount]
    )  # 植物库存

    # Plant names and prices
    plants_names = [
        f"{PLANTS_TIERS[row['tier']]} {PLANTS_LABLES[row['species']]}"
        for index, row in df.iterrows()
    ]
    price = df[currency]  # 植物单价
    sold_prices = np.array(price * (1 + extra_rate))

    # Initialize the master problem
    model = Model("BewilderingBlooms")

    # Decision variables in master problem
    x = [
        model.addVar(
            vtype="I", name=f"x_{i}", lb=0, ub=int(stocks[i]) if stocks[i] else 0
        )
        for i in range(len(stocks))
    ]

    obj1 = quicksum(sold_prices[i] * x[i] for i in range(len(stocks)))
    obj2 = quicksum(x[i] for i in range(len(stocks)))

    # Objective: maximize total value of sold plants
    model.setObjective(obj1, "maximize")

    model.addCons(obj1 <= budget)

    # first optimize
    model.hideOutput()
    model.optimize()

    if model.getStatus() == "optimal":
        optimal_total_value = model.getObjVal()
        print(
            f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] First-stage optimal total value: {optimal_total_value}"
        )
        model.freeTransform()

        model.setObjective(
            obj2, "maximize" if strategy == "MinimizeStock" else "minimize"
        )
        model.addCons(obj1 == optimal_total_value)
        model.optimize()

        # Final solution processing
        solution = []
        total_price = 0
        total_count = 0

        if model.getStatus() == "optimal":
            for i, var in enumerate(x):
                if (v := round(model.getVal(var))) > 0 and sold_prices[i] > 0:
                    solution.append(
                        f"{plants_names[i]} ({sold_prices[i]} {currency}): {v}\n"
                    )
                    total_price += v * sold_prices[i]
                    total_count += v

            if optimal_total_value == budget:
                return f"Great! Found a combination of items with a total value equal to the budget ({budget} {currency}).😃\n\n{''.join(solution)}\nTotal value: {total_price} {currency}\nTotal count: {total_count}"  # Count: {int(model.getObjVal())}

            return f"Oops! {round(budget - optimal_total_value)} {currency} short of the target value ({budget} {currency}).😣\n\n{''.join(solution)}\nTotal value: {total_price} {currency}\nTotal count: {total_count}"  # Count: {int(model.getObjVal())}

        return "No solution found for the second optimization!"

    return "No solution found for the first optimization!"


# 高亮每种植物的最高品质
css = """
.highlight-first-gold {background-color: #fafad2}
.highlight-first-gems {background-color: #fed9b4}
"""


def show_checkboxgroup(currency, select_all=False):
    """
    根据选定的货币显示选择框。
    """
    plants_tuples: list[tuple[str, str]] = [
        (PLANTS_LABLES[pl], pl)
        for pl in PLANTS_LABLES.keys()
        if check_currency(pl, currency)
    ]

    if select_all:
        default_value = [v for (n, v) in plants_tuples if select_all]
    else:
        default_value = None

    checkbox_group_component = gr.CheckboxGroup(
        choices=plants_tuples,
        value=default_value,
        type="value",
        label="Plants",
        info="Select plants",
        interactive=True,
    )
    return checkbox_group_component


def show_plant_boxes(currency, plants=None):
    _inventory = {}
    species_set = set()
    species_count = 0
    new_species = False

    for _, row in df.iterrows():
        if (
            row[currency] != 0
            and row["tier"] != "feeble"
            and (not plants or row["species"] in plants)
        ):
            species_set.add(row["species"])
            inventory_key = f"{row['species']}_{row['tier']}"
            _inventory[inventory_key] = gr.Number(
                label=PLANTS_LABLES[row["species"]],
                info=f"{PLANTS_TIERS[row['tier']]} ${row[currency]}",
                value=0,
                precision=0,
                minimum=0,
                maximum=500,
                step=10,
                visible=True,
                elem_classes=(
                    f"highlight-first-{currency}"
                    if len(species_set) > species_count
                    else None
                ),
            )
            species_count = len(species_set)
        else:
            _inventory[f"{row['species']}_{row['tier']}"] = gr.Number(
                value=0, visible=False
            )

    return list(_inventory.values())


def handle_currency(currency):
    """
    根据选定的货币类型更新库存组件"""
    return [False, show_checkboxgroup(currency)] + show_plant_boxes(currency, None)


def handle_select_all(initial_state, currency):
    return [(not initial_state)] + [show_checkboxgroup(currency, not initial_state)]


with gr.Blocks(css=css) as demo:
    gr.Markdown(
        """
        <center><font size=8>HPMA Bewildering Blooms Calculator👨🏻‍🌾</font></center>

        This program is essentially a solver for a variant of the knapsack problem. 
        Another more versatile [application](https://huggingface.co/spaces/oh-my-dear-ai/easy-knapsack-problem).
        """
    )

    # Create a Gradio interface with a column layout
    with gr.Column():
        # Add a row for the currency selection
        currency_radio = gr.Radio(
            choices=["gold", "gems"],
            value="gold",
            type="value",
            label="Currency",
            info="Select the currency:",
            render=True,
        )
        # Add a row for the budget input
        budget = gr.Number(
            label="Target",
            info="Gabby's Budget:",  # "葭碧の金币预算：",
            value=0,
            minimum=0,
            maximum=20000,
            step=100,
        )
        acquisition_rate = gr.Dropdown(
            choices=[
                "0(Gabby's Acquisition)",
                "+100%(HVA for Budding & Novice)",
                "+200%(HVA for Junior & Practiced)",
                "+300%(HVA for Natural & Master)",
            ],
            value="0(Gabby's Acquisition)",
            type="index",
            label="Extra Acquisition Rate",
            info="Select your high-value acquisition rate:",
        )

        # Add a radio selection for the strategy
        selected_strategy = gr.Radio(
            [
                (
                    "Minimize the number of plants sold (prioritize high-priced plants)",
                    "MaximizeStock",
                ),
                (
                    "Maximize the number of plants sold (prioritize low-priced plants)",
                    "MinimizeStock",
                ),
            ],
            value="MaximizeStock",
            label="Strategies",
            info="Select a strategy:",
        )

        plants_filter = show_checkboxgroup(currency_radio.value)

        with gr.Row():
            select_all_state = gr.State(False)
            select_all_button = gr.Button(value="Select All⭕", size="sm")
            filter_button = gr.Button(value="Filter Plants🔍", size="lg")

    # Create the dynamic plant inventory inputs
    with gr.Row() as inventory_row:
        inventory = show_plant_boxes(currency_radio.value, plants_filter.value)

    # Add a row for the Clear and Calculate buttons
    with gr.Row():
        inventory_clear_btn = gr.ClearButton(inventory, size="sm", value="❌Clear")

        # Add a button to trigger the calculation
        inventory_submit_btn = gr.Button(value="🛠Calculate")

    # Add a row for the result textbox
    with gr.Row():
        result = gr.Textbox(label="Output")

    # Set up the button click event to call the calculator function
    inventory_submit_btn.click(
        calculator,
        inputs=[currency_radio, budget, selected_strategy, acquisition_rate]
        + inventory,
        outputs=[result],
        api_name=False,
    )

    # Update the inventory when the currency changes
    currency_radio.change(
        fn=handle_currency,  # Adjusted function to return only the components
        inputs=[currency_radio],
        outputs=[
            select_all_state,
            plants_filter,
        ]
        + inventory,  # Update each child in the inventory_row
    )

    filter_button.click(
        fn=show_plant_boxes,
        inputs=[currency_radio, plants_filter],
        outputs=inventory,
    )

    select_all_button.click(
        fn=handle_select_all,
        inputs=[select_all_state, currency_radio],
        outputs=[select_all_state, plants_filter],
    )

# Launch the Gradio application
demo.queue(api_open=False)
demo.launch(max_threads=5, share=False)