import os
import matplotlib.pyplot as plt

class BacktestCharts:
    def __init__(self, charts_dir="charts"):
        self.charts_dir = charts_dir
        os.makedirs(self.charts_dir, exist_ok=True)

    def plot_profit_ratio_vs_stop_loss(self, results, filename="profit_ratio_vs_stop_loss.png"):
        """
        Plots profit ratio vs stop loss percentage for each timeframe.

        Parameters:
        - results: list of dicts, each with keys: 'timeframe', 'stop_loss_pct', 'profit_ratio'
        - filename: output filename (will be saved in charts_dir)
        """
        # Organize data by timeframe
        from collections import defaultdict
        data = defaultdict(lambda: {"stop_loss_pct": [], "profit_ratio": []})
        for row in results:
            tf = row["timeframe"]
            data[tf]["stop_loss_pct"].append(row["stop_loss_pct"])
            data[tf]["profit_ratio"].append(row["profit_ratio"])

        plt.figure(figsize=(10, 6))
        for tf, vals in data.items():
            # Sort by stop_loss_pct for smooth lines
            sorted_pairs = sorted(zip(vals["stop_loss_pct"], vals["profit_ratio"]))
            stop_loss, profit_ratio = zip(*sorted_pairs)
            plt.plot(
                [s * 100 for s in stop_loss],  # Convert to percent
                profit_ratio,
                marker="o",
                label=tf
            )

        plt.xlabel("Stop Loss (%)")
        plt.ylabel("Profit Ratio")
        plt.title("Profit Ratio vs Stop Loss (%) per Timeframe")
        plt.legend(title="Timeframe")
        plt.grid(True, linestyle="--", alpha=0.5)
        plt.tight_layout()

        output_path = os.path.join(self.charts_dir, filename)
        plt.savefig(output_path)
        plt.close()

    def plot_average_trade_vs_stop_loss(self, results, filename="average_trade_vs_stop_loss.png"):
        """
        Plots average trade vs stop loss percentage for each timeframe.

        Parameters:
        - results: list of dicts, each with keys: 'timeframe', 'stop_loss_pct', 'average_trade'
        - filename: output filename (will be saved in charts_dir)
        """
        from collections import defaultdict
        data = defaultdict(lambda: {"stop_loss_pct": [], "average_trade": []})
        for row in results:
            tf = row["timeframe"]
            if "average_trade" not in row:
                continue  # Skip rows without average_trade
            data[tf]["stop_loss_pct"].append(row["stop_loss_pct"])
            data[tf]["average_trade"].append(row["average_trade"])

        plt.figure(figsize=(10, 6))
        for tf, vals in data.items():
            # Sort by stop_loss_pct for smooth lines
            sorted_pairs = sorted(zip(vals["stop_loss_pct"], vals["average_trade"]))
            stop_loss, average_trade = zip(*sorted_pairs)
            plt.plot(
                [s * 100 for s in stop_loss],  # Convert to percent
                average_trade,
                marker="o",
                label=tf
            )

        plt.xlabel("Stop Loss (%)")
        plt.ylabel("Average Trade")
        plt.title("Average Trade vs Stop Loss (%) per Timeframe")
        plt.legend(title="Timeframe")
        plt.grid(True, linestyle="--", alpha=0.5)
        plt.tight_layout()

        output_path = os.path.join(self.charts_dir, filename)
        plt.savefig(output_path)
        plt.close()