import random
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import gaussian_kde
import tkinter as tk
from tkinter import ttk, messagebox
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import json
from datetime import datetime
import os

def trading_simulation(initial_capital, investment_percent, iterations, returns, probs):
    """Simulate a single person's trading strategy with custom returns and probabilities"""
    capital = initial_capital
    capital_history = [capital]

    total_prob = sum(probs)
    if total_prob != 1:
        probs = [p / total_prob for p in probs]

    for _ in range(iterations):
        investment = capital * investment_percent
        outcome = random.choices(range(len(returns)), weights=probs, k=1)[0]
        capital += investment * (returns[outcome] - 1)
        capital_history.append(capital)

    return capital_history


def simulate_multiple_people(initial_capital, investment_percent, people, iterations, returns, probs):
    """Simulate multiple people trading"""
    return [trading_simulation(initial_capital, investment_percent, iterations, returns, probs) for _ in range(people)]


def create_plots(all_histories, investment_percent, iterations, initial_capital):
    """Create evolution and distribution plots with larger size"""
    fig1, ax1 = plt.subplots(figsize=(8, 5))
    for history in all_histories:
        ax1.plot(history, alpha=0.3)
    ax1.set_title(f'Capital Evolution ({len(all_histories)} People, {investment_percent * 100}%, {iterations} Trades)')
    ax1.set_xlabel('Number of Trades')
    ax1.set_ylabel('Capital')
    ax1.grid(True)

    final_capitals = np.array([history[-1] for history in all_histories])
    capital_changes = ((final_capitals - initial_capital) / initial_capital) * 100
    log_changes = np.where(
        capital_changes >= 0,
        np.log10(1 + capital_changes / 100),
        -np.log10(1 + np.abs(capital_changes) / 100)
    )
    kde = gaussian_kde(log_changes)
    x_min, x_max = np.percentile(log_changes, [5, 95])
    if x_max - x_min < 0.1:
        x_min, x_max = min(log_changes, default=-2), max(log_changes, default=2)
    x_range = np.linspace(x_min, x_max, 200)
    density = kde(x_range)

    fig2, ax2 = plt.subplots(figsize=(8, 5))
    ax2.plot(x_range, density, color='blue', label='Probability Density')
    ax2.fill_between(x_range, 0, density, where=(x_range < 0), color='red', alpha=0.3, label='Loss')
    ax2.fill_between(x_range, 0, density, where=(x_range > 0), color='green', alpha=0.3, label='Profit')
    ax2.axvline(x=0, color='black', linestyle='--', alpha=0.5, label='No Change (0%)')
    ax2.set_title(f'Final Capital Change Distribution ({investment_percent * 100}%)')
    ax2.set_xlabel('Log of % Change')
    ax2.set_ylabel('Density')
    ax2.grid(True)
    ax2.legend()
    ax2.minorticks_on()
    ax2.grid(which='minor', alpha=0.2)

    return fig1, fig2


def run_simulation():
    """Run the simulation with user inputs"""
    try:
        initial_capital = float(entry_initial_capital.get())
        investment_percent = float(entry_investment_percent.get()) / 100
        people = int(entry_people.get())
        iterations = int(entry_iterations.get())

        returns = []
        probs = []
        for return_entry, prob_entry in probability_entries:
            return_val = return_entry.get().strip()
            prob_val = prob_entry.get().strip()
            if not return_val or not prob_val:
                raise ValueError("All return and probability fields must be filled.")
            returns.append(float(return_val))
            probs.append(float(prob_val) / 100)

        # 只检查概率是否为负数，不限制return
        if any(p < 0 for p in probs):
            raise ValueError("Probabilities must be non-negative.")
        if initial_capital < 0 or investment_percent < 0 or people <= 0 or iterations <= 0:
            raise ValueError("Initial capital, investment %, people, and iterations must be positive.")

        all_histories = simulate_multiple_people(initial_capital, investment_percent, people, iterations, returns,
                                                 probs)
        final_capitals = np.array([history[-1] for history in all_histories])

        result_text.delete(1.0, tk.END)
        result_text.insert(tk.END, f"Initial Capital: {initial_capital:.2f}\n")
        result_text.insert(tk.END, f"Investment Percent: {investment_percent * 100:.2f}%\n")
        result_text.insert(tk.END, f"Number of People: {people}\n")
        result_text.insert(tk.END, f"Number of Trades: {iterations}\n")
        result_text.insert(tk.END, f"Average Final Capital: {np.mean(final_capitals):.2f}\n")
        result_text.insert(tk.END, f"Median Final Capital: {np.median(final_capitals):.2f}\n")
        result_text.insert(tk.END, f"Max Final Capital: {np.max(final_capitals):.2f}\n")
        result_text.insert(tk.END, f"Min Final Capital: {np.min(final_capitals):.2f}\n")
        result_text.insert(tk.END,
                           f"Average Capital Change: {((np.mean(final_capitals) - initial_capital) / initial_capital * 100):.2f}%\n")

        if save_to_history.get():  # Only save if checkbox is checked
            run_data = {
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "initial_capital": initial_capital,
                "investment_percent": investment_percent * 100,
                "people": people,
                "iterations": iterations,
                "returns": returns,
                "probs": [p * 100 for p in probs],
                "avg_final_capital": float(np.mean(final_capitals))
            }
            run_history.append(run_data)
            save_history_to_file()
            update_history_dropdown()

        for widget in tab1.winfo_children():
            widget.destroy()
        for widget in tab2.winfo_children():
            widget.destroy()

        fig1, fig2 = create_plots(all_histories, investment_percent, iterations, initial_capital)

        canvas1 = FigureCanvasTkAgg(fig1, master=tab1)
        canvas1.draw()
        canvas1.get_tk_widget().pack(fill=tk.BOTH, expand=1)

        canvas2 = FigureCanvasTkAgg(fig2, master=tab2)
        canvas2.draw()
        canvas2.get_tk_widget().pack(fill=tk.BOTH, expand=1)

    except ValueError as e:
        messagebox.showerror("Input Error", f"Invalid input: {e}")


def add_probability_pair(return_val="", prob_val=""):
    """Add a new return-probability pair with optional default values"""
    row = len(probability_entries) + 3
    return_label = tk.Label(inner_input_frame, text=f"Return {len(probability_entries) + 1}:")
    return_label.grid(row=row, column=0, padx=5, pady=2, sticky="e")
    return_entry = tk.Entry(inner_input_frame, width=10)
    return_entry.insert(0, return_val)
    return_entry.grid(row=row, column=1, padx=5, pady=2)

    prob_label = tk.Label(inner_input_frame, text=f"Prob {len(probability_entries) + 1} (%):")
    prob_label.grid(row=row, column=2, padx=5, pady=2, sticky="e")
    prob_entry = tk.Entry(inner_input_frame, width=10)
    prob_entry.insert(0, prob_val)
    prob_entry.grid(row=row, column=3, padx=5, pady=2)

    probability_entries.append((return_entry, prob_entry))
    update_input_scrollbar()


def remove_probability_pair():
    """Remove the last return-probability pair"""
    if probability_entries:
        last_return_entry, last_prob_entry = probability_entries.pop()
        last_return_entry.master.grid_forget()  # Remove label
        last_prob_entry.master.grid_forget()  # Remove label
        last_return_entry.destroy()
        last_prob_entry.destroy()
        update_input_scrollbar()


def update_input_scrollbar():
    """Update the scrollbar for input frame"""
    inner_input_frame.update_idletasks()
    input_canvas.configure(scrollregion=input_canvas.bbox("all"))


def save_history_to_file():
    """Save run history to a text file"""
    with open("cache/run_history.txt", "w") as f:
        json.dump(run_history, f, indent=4)


def load_history_from_file():
    """Load run history from a text file"""
    global run_history
    if os.path.exists("cache/run_history.txt"):
        with open("cache/run_history.txt", "r") as f:
            run_history = json.load(f)
    else:
        run_history = []


def update_history_dropdown():
    """Update the history dropdown with saved runs"""
    history_menu['menu'].delete(0, 'end')
    for i, run in enumerate(run_history):
        label = f"{run['timestamp']} - Avg: {run['avg_final_capital']:.2f}"
        history_menu['menu'].add_command(label=label, command=lambda idx=i: load_history(idx))


def load_history(index):
    """Load a previous run's parameters"""
    run = run_history[index]
    entry_initial_capital.delete(0, tk.END)
    entry_initial_capital.insert(0, str(run['initial_capital']))
    entry_investment_percent.delete(0, tk.END)
    entry_investment_percent.insert(0, str(run['investment_percent']))
    entry_people.delete(0, tk.END)
    entry_people.insert(0, str(run['people']))
    entry_iterations.delete(0, tk.END)
    entry_iterations.insert(0, str(run['iterations']))

    for return_entry, prob_entry in probability_entries:
        return_entry.master.grid_forget()
        prob_entry.master.grid_forget()
        return_entry.destroy()
        prob_entry.destroy()
    probability_entries.clear()

    for r, p in zip(run['returns'], run['probs']):
        add_probability_pair(str(r), str(p))


# Create GUI window
root = tk.Tk()
root.title("Trading Simulation with Custom Returns")
root.geometry("1200x900")

window_height = 900
top_height = int(window_height * 0.3)
plot_height = int(window_height * 0.7)

# Load history at startup
run_history = []
load_history_from_file()

# Top frame
top_frame = tk.Frame(root, height=top_height)
top_frame.pack(side=tk.TOP, fill=tk.X, padx=10, pady=10)
top_frame.pack_propagate(False)

# Input frame with scrollbar
input_frame = tk.LabelFrame(top_frame, text="Simulation Parameters", padx=10, pady=10)
input_frame.pack(side=tk.LEFT, fill=tk.Y, expand=1)

input_canvas = tk.Canvas(input_frame)
input_scrollbar = tk.Scrollbar(input_frame, orient="vertical", command=input_canvas.yview)
inner_input_frame = tk.Frame(input_canvas)
input_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
input_canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=1)
input_canvas.create_window((0, 0), window=inner_input_frame, anchor="nw")
input_canvas.configure(yscrollcommand=input_scrollbar.set)

tk.Label(inner_input_frame, text="Initial Capital:").grid(row=0, column=0, padx=5, pady=5, sticky="e")
entry_initial_capital = tk.Entry(inner_input_frame, width=15)
entry_initial_capital.insert(0, "10000.0")
entry_initial_capital.grid(row=0, column=1, padx=5, pady=5)

tk.Label(inner_input_frame, text="Investment %:").grid(row=0, column=2, padx=5, pady=5, sticky="e")
entry_investment_percent = tk.Entry(inner_input_frame, width=15)
entry_investment_percent.insert(0, "10.0")
entry_investment_percent.grid(row=0, column=3, padx=5, pady=5)

tk.Label(inner_input_frame, text="People:").grid(row=1, column=0, padx=5, pady=5, sticky="e")
entry_people = tk.Entry(inner_input_frame, width=15)
entry_people.insert(0, "100")
entry_people.grid(row=1, column=1, padx=5, pady=5)

tk.Label(inner_input_frame, text="Trades:").grid(row=1, column=2, padx=5, pady=5, sticky="e")
entry_iterations = tk.Entry(inner_input_frame, width=15)
entry_iterations.insert(0, "100")
entry_iterations.grid(row=1, column=3, padx=5, pady=5)

run_button = tk.Button(inner_input_frame, text="Run Simulation", command=run_simulation)
run_button.grid(row=2, column=0, columnspan=2, pady=5)

add_button = tk.Button(inner_input_frame, text="+", command=lambda: add_probability_pair())
add_button.grid(row=2, column=3, padx=5, pady=5)

remove_button = tk.Button(inner_input_frame, text="-", command=remove_probability_pair)
remove_button.grid(row=2, column=2, padx=5, pady=5)

default_returns = [20.0, 10.0, 5.0, 2.0, 1.5, 1.2, 1.05, 1.0, 0.9, 0.5, 0.4, 0.2, 0.1, -10.0]
default_probs = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 20.0, 20.0, 9.0, 15.0, 15.0, 10.0]

probability_entries = []
for r, p in zip(default_returns, default_probs):
    add_probability_pair(str(r), str(p))

inner_input_frame.bind("<Configure>", lambda e: update_input_scrollbar())

# Result frame
result_frame = tk.LabelFrame(top_frame, text="Results", padx=10, pady=10)
result_frame.pack(side=tk.RIGHT, fill=tk.Y, expand=1)

result_text = tk.Text(result_frame, height=10, width=40)
result_text.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
result_scrollbar = tk.Scrollbar(result_frame, orient="vertical", command=result_text.yview)
result_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
result_text.configure(yscrollcommand=result_scrollbar.set)

history_frame = tk.Frame(result_frame)
history_frame.pack(side=tk.BOTTOM, fill=tk.X, pady=5)
tk.Label(history_frame, text="Run History:").pack(side=tk.LEFT, padx=5)
history_var = tk.StringVar()
history_menu = tk.OptionMenu(history_frame, history_var, "No runs yet")
history_menu.pack(side=tk.LEFT, fill=tk.X, expand=1)

# Checkbox for saving to history
save_to_history = tk.BooleanVar(value=True)  # Default checked
save_checkbox = tk.Checkbutton(history_frame, text="Save to History", variable=save_to_history)
save_checkbox.pack(side=tk.LEFT, padx=5)

# Populate history dropdown with loaded data
update_history_dropdown()

# Plot frame with tabs
plot_outer_frame = tk.LabelFrame(root, text="Plots", padx=10, pady=10, height=plot_height)
plot_outer_frame.pack(fill=tk.BOTH, expand=1)
plot_outer_frame.pack_propagate(False)

notebook = ttk.Notebook(plot_outer_frame)
notebook.pack(fill=tk.BOTH, expand=1)

tab1 = ttk.Frame(notebook)
tab2 = ttk.Frame(notebook)
notebook.add(tab1, text="Capital Evolution")
notebook.add(tab2, text="Distribution")

root.mainloop()