import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


# ============================================================
# 設定
# ============================================================

INPUT_FILE = "results/trade_base.csv"
OUTPUT_DIR = "results"

os.makedirs(
    OUTPUT_DIR,
    exist_ok=True
)

INITIAL_CAPITAL = 700000


CANDIDATES = {
    "BASE": None,
    "A_0.44_0.88": (0.44, 0.88),
    "B_0.44_1.00": (0.44, 1.00),
    "C_0.50_1.00": (0.50, 1.00),
    "D_0.60_1.00": (0.60, 1.00),
}


# ============================================================
# 出力ファイル
# ============================================================

SUMMARY_CSV = os.path.join(
    OUTPUT_DIR,
    "equity_drawdown_summary.csv"
)

EQUITY_CSV = os.path.join(
    OUTPUT_DIR,
    "equity_curve_data.csv"
)

DRAWDOWN_CSV = os.path.join(
    OUTPUT_DIR,
    "drawdown_curve_data.csv"
)

EQUITY_GRAPH = os.path.join(
    OUTPUT_DIR,
    "equity_curve_comparison.png"
)

DRAWDOWN_GRAPH = os.path.join(
    OUTPUT_DIR,
    "drawdown_curve_comparison.png"
)

MAX_DD_GRAPH = os.path.join(
    OUTPUT_DIR,
    "max_drawdown_comparison.png"
)


# ============================================================
# データ読み込み
# ============================================================

df = pd.read_csv(
    INPUT_FILE
)

print(
    "========================================"
)

print(
    "RELATIVE BAND EQUITY / DRAWDOWN ANALYSIS"
)

print(
    "========================================"
)

print(
    f"\n読み込み件数: {len(df)}"
)


# ============================================================
# 必要列確認
# ============================================================

required_columns = [
    "entry_time",
    "entry_band_distance_pips",
    "entry_atr14_pips",
    "profit",
]

missing_columns = [
    column
    for column in required_columns
    if column not in df.columns
]

if missing_columns:
    raise ValueError(
        "必要な列がありません: "
        + ", ".join(missing_columns)
    )


# ============================================================
# 型変換
# ============================================================

df["entry_time"] = pd.to_datetime(
    df["entry_time"],
    errors="coerce"
)

numeric_columns = [
    "entry_band_distance_pips",
    "entry_atr14_pips",
    "profit",
]

for column in numeric_columns:

    df[column] = pd.to_numeric(
        df[column],
        errors="coerce"
    )


# ============================================================
# 分析対象データ
# ============================================================

analysis_df = df[
    df["entry_time"].notna()
    &
    df["entry_band_distance_pips"].notna()
    &
    df["entry_atr14_pips"].notna()
    &
    df["profit"].notna()
    &
    (
        df["entry_atr14_pips"] > 0
    )
].copy()

if analysis_df.empty:
    raise ValueError(
        "分析対象データがありません。"
    )


# ============================================================
# Relative Band Distance
# ============================================================

analysis_df[
    "relative_band_distance"
] = (
    analysis_df[
        "entry_band_distance_pips"
    ]
    /
    analysis_df[
        "entry_atr14_pips"
    ]
)

analysis_df[
    "relative_band_distance"
] = (
    analysis_df[
        "relative_band_distance"
    ]
    .replace(
        [
            np.inf,
            -np.inf
        ],
        np.nan
    )
)

analysis_df = analysis_df[
    analysis_df[
        "relative_band_distance"
    ].notna()
].copy()


# ============================================================
# 時系列順
# ============================================================

analysis_df = (
    analysis_df
    .sort_values(
        "entry_time"
    )
    .reset_index(
        drop=True
    )
)

print(
    f"分析対象件数: {len(analysis_df)}"
)

print(
    f"分析開始日: "
    f"{analysis_df['entry_time'].min()}"
)

print(
    f"分析終了日: "
    f"{analysis_df['entry_time'].max()}"
)


# ============================================================
# PF
# ============================================================

def calculate_pf(profits):

    gross_profit = profits[
        profits > 0
    ].sum()

    gross_loss = -profits[
        profits < 0
    ].sum()

    if gross_loss == 0:

        if gross_profit > 0:
            return np.inf

        return np.nan

    return (
        gross_profit
        /
        gross_loss
    )


# ============================================================
# 条件抽出
# ============================================================

def filter_candidate(
    source_df,
    condition
):

    if condition is None:

        return (
            source_df
            .copy()
            .reset_index(drop=True)
        )

    lower, upper = condition

    filtered = source_df[
        (
            source_df[
                "relative_band_distance"
            ]
            >= lower
        )
        &
        (
            source_df[
                "relative_band_distance"
            ]
            <= upper
        )
    ].copy()

    return (
        filtered
        .sort_values(
            "entry_time"
        )
        .reset_index(
            drop=True
        )
    )


# ============================================================
# 資産曲線・DD計算
# ============================================================

def analyze_equity(
    candidate_df,
    initial_capital
):

    result = (
        candidate_df
        .copy()
        .reset_index(drop=True)
    )

    result[
        "trade_number"
    ] = (
        np.arange(
            len(result)
        )
        + 1
    )

    result[
        "cumulative_profit"
    ] = (
        result[
            "profit"
        ]
        .cumsum()
    )

    result[
        "equity"
    ] = (
        initial_capital
        +
        result[
            "cumulative_profit"
        ]
    )

    # 初期資金も過去最高値として扱う
    running_max_values = []

    running_max = initial_capital

    for equity in result["equity"]:

        running_max = max(
            running_max,
            equity
        )

        running_max_values.append(
            running_max
        )

    result[
        "running_max"
    ] = running_max_values

    result[
        "drawdown_amount"
    ] = (
        result[
            "equity"
        ]
        -
        result[
            "running_max"
        ]
    )

    result[
        "drawdown_pct"
    ] = (
        result[
            "drawdown_amount"
        ]
        /
        result[
            "running_max"
        ]
        * 100
    )

    return result


# ============================================================
# 最大DD情報
# ============================================================

def get_drawdown_info(
    equity_df,
    initial_capital
):

    if equity_df.empty:

        return {
            "max_drawdown_amount": np.nan,
            "max_drawdown_pct": np.nan,
            "dd_peak_time": pd.NaT,
            "dd_bottom_time": pd.NaT,
            "recovery_time": pd.NaT,
            "recovery_trades": np.nan,
        }

    bottom_index = (
        equity_df[
            "drawdown_pct"
        ]
        .idxmin()
    )

    bottom_row = (
        equity_df
        .loc[
            bottom_index
        ]
    )

    max_drawdown_amount = abs(
        bottom_row[
            "drawdown_amount"
        ]
    )

    max_drawdown_pct = abs(
        bottom_row[
            "drawdown_pct"
        ]
    )

    peak_equity = (
        bottom_row[
            "running_max"
        ]
    )

    before_bottom = (
        equity_df
        .loc[
            :bottom_index
        ]
    )

    peak_candidates = before_bottom[
        before_bottom[
            "equity"
        ]
        >= peak_equity
    ]

    if peak_candidates.empty:

        peak_time = pd.NaT
        peak_trade_number = 0

    else:

        peak_row = (
            peak_candidates
            .iloc[-1]
        )

        peak_time = (
            peak_row[
                "entry_time"
            ]
        )

        peak_trade_number = int(
            peak_row[
                "trade_number"
            ]
        )

    bottom_time = (
        bottom_row[
            "entry_time"
        ]
    )

    after_bottom = (
        equity_df
        .loc[
            bottom_index + 1:
        ]
    )

    recovery_candidates = after_bottom[
        after_bottom[
            "equity"
        ]
        >= peak_equity
    ]

    if recovery_candidates.empty:

        recovery_time = pd.NaT
        recovery_trades = np.nan

    else:

        recovery_row = (
            recovery_candidates
            .iloc[0]
        )

        recovery_time = (
            recovery_row[
                "entry_time"
            ]
        )

        recovery_trades = (
            int(
                recovery_row[
                    "trade_number"
                ]
            )
            -
            peak_trade_number
        )

    return {
        "max_drawdown_amount":
            max_drawdown_amount,

        "max_drawdown_pct":
            max_drawdown_pct,

        "dd_peak_time":
            peak_time,

        "dd_bottom_time":
            bottom_time,

        "recovery_time":
            recovery_time,

        "recovery_trades":
            recovery_trades,
    }


# ============================================================
# 分析
# ============================================================

summary_rows = []
equity_rows = []
drawdown_rows = []

equity_results = {}


for candidate_name, condition in CANDIDATES.items():

    candidate_df = filter_candidate(
        analysis_df,
        condition
    )

    if candidate_df.empty:

        print(
            f"\n{candidate_name}: "
            "対象取引なし"
        )

        continue

    equity_df = analyze_equity(
        candidate_df,
        INITIAL_CAPITAL
    )

    equity_results[
        candidate_name
    ] = equity_df

    profits = (
        candidate_df[
            "profit"
        ]
    )

    trades = len(
        candidate_df
    )

    wins = int(
        (
            profits > 0
        ).sum()
    )

    win_rate = (
        wins
        /
        trades
        *
        100
    )

    total_profit = (
        profits.sum()
    )

    average_profit = (
        profits.mean()
    )

    gross_profit = (
        profits[
            profits > 0
        ].sum()
    )

    gross_loss = -(
        profits[
            profits < 0
        ].sum()
    )

    profit_factor = (
        calculate_pf(
            profits
        )
    )

    final_equity = (
        equity_df[
            "equity"
        ]
        .iloc[-1]
    )

    dd_info = (
        get_drawdown_info(
            equity_df,
            INITIAL_CAPITAL
        )
    )

    summary_rows.append(
        {
            "candidate":
                candidate_name,

            "trades":
                trades,

            "wins":
                wins,

            "win_rate":
                win_rate,

            "total_profit":
                total_profit,

            "average_profit":
                average_profit,

            "gross_profit":
                gross_profit,

            "gross_loss":
                gross_loss,

            "profit_factor":
                profit_factor,

            "initial_capital":
                INITIAL_CAPITAL,

            "final_equity":
                final_equity,

            "max_drawdown_amount":
                dd_info[
                    "max_drawdown_amount"
                ],

            "max_drawdown_pct":
                dd_info[
                    "max_drawdown_pct"
                ],

            "dd_peak_time":
                dd_info[
                    "dd_peak_time"
                ],

            "dd_bottom_time":
                dd_info[
                    "dd_bottom_time"
                ],

            "recovery_time":
                dd_info[
                    "recovery_time"
                ],

            "recovery_trades":
                dd_info[
                    "recovery_trades"
                ],
        }
    )

    temp_equity = equity_df[
        [
            "entry_time",
            "trade_number",
            "profit",
            "cumulative_profit",
            "equity",
        ]
    ].copy()

    temp_equity[
        "candidate"
    ] = candidate_name

    equity_rows.append(
        temp_equity
    )

    temp_drawdown = equity_df[
        [
            "entry_time",
            "trade_number",
            "equity",
            "running_max",
            "drawdown_amount",
            "drawdown_pct",
        ]
    ].copy()

    temp_drawdown[
        "candidate"
    ] = candidate_name

    drawdown_rows.append(
        temp_drawdown
    )


# ============================================================
# DataFrame化
# ============================================================

summary_df = pd.DataFrame(
    summary_rows
)

equity_output_df = pd.concat(
    equity_rows,
    ignore_index=True
)

drawdown_output_df = pd.concat(
    drawdown_rows,
    ignore_index=True
)


# ============================================================
# CSV保存
# ============================================================

summary_df.to_csv(
    SUMMARY_CSV,
    index=False
)

equity_output_df.to_csv(
    EQUITY_CSV,
    index=False
)

drawdown_output_df.to_csv(
    DRAWDOWN_CSV,
    index=False
)


# ============================================================
# 資産曲線
# ============================================================

plt.figure(
    figsize=(12, 7)
)

for (
    candidate_name,
    equity_df
) in equity_results.items():

    plt.plot(
        equity_df[
            "entry_time"
        ],
        equity_df[
            "equity"
        ],
        label=candidate_name
    )

plt.axhline(
    INITIAL_CAPITAL,
    linestyle="--",
    linewidth=1
)

plt.title(
    "Equity Curve Comparison"
)

plt.xlabel(
    "Entry Time"
)

plt.ylabel(
    "Equity"
)

plt.legend()

plt.grid(
    alpha=0.3
)

plt.tight_layout()

plt.savefig(
    EQUITY_GRAPH,
    dpi=150
)

plt.close()


# ============================================================
# DD曲線
# ============================================================

plt.figure(
    figsize=(12, 7)
)

for (
    candidate_name,
    equity_df
) in equity_results.items():

    plt.plot(
        equity_df[
            "entry_time"
        ],
        equity_df[
            "drawdown_pct"
        ],
        label=candidate_name
    )

plt.axhline(
    0,
    linestyle="--",
    linewidth=1
)

plt.title(
    "Drawdown Curve Comparison"
)

plt.xlabel(
    "Entry Time"
)

plt.ylabel(
    "Drawdown (%)"
)

plt.legend()

plt.grid(
    alpha=0.3
)

plt.tight_layout()

plt.savefig(
    DRAWDOWN_GRAPH,
    dpi=150
)

plt.close()


# ============================================================
# 最大DD比較グラフ
# ============================================================

plot_df = (
    summary_df
    .sort_values(
        "max_drawdown_pct"
    )
    .copy()
)

plt.figure(
    figsize=(10, 6)
)

plt.bar(
    plot_df[
        "candidate"
    ],
    plot_df[
        "max_drawdown_pct"
    ]
)

plt.title(
    "Maximum Drawdown Comparison"
)

plt.xlabel(
    "Candidate"
)

plt.ylabel(
    "Maximum Drawdown (%)"
)

plt.xticks(
    rotation=20
)

plt.grid(
    axis="y",
    alpha=0.3
)

plt.tight_layout()

plt.savefig(
    MAX_DD_GRAPH,
    dpi=150
)

plt.close()


# ============================================================
# 結果表示
# ============================================================

print(
    "\n========================================"
)

print(
    "SUMMARY"
)

print(
    "========================================"
)

display_columns = [
    "candidate",
    "trades",
    "win_rate",
    "total_profit",
    "average_profit",
    "profit_factor",
    "final_equity",
    "max_drawdown_amount",
    "max_drawdown_pct",
    "dd_peak_time",
    "dd_bottom_time",
    "recovery_time",
    "recovery_trades",
]

print(
    summary_df[
        display_columns
    ].to_string(
        index=False
    )
)


print(
    "\n========================================"
)

print(
    "出力ファイル"
)

print(
    "========================================"
)

print(
    SUMMARY_CSV
)

print(
    EQUITY_CSV
)

print(
    DRAWDOWN_CSV
)

print(
    EQUITY_GRAPH
)

print(
    DRAWDOWN_GRAPH
)

print(
    MAX_DD_GRAPH
)