from datetime import date, timedelta
from pathlib import Path

from iching.utils import bz
from iching.utils import bzshagod as sg


def day_info(d: date):
    b = bz.getCalendar10kGodEarthStem(d.year, d.month, d.day, 22, 59, 59)
    g = b["day"]["g"]
    e = b["day"]["e"]
    return g, e, b


def next_dates_with_day_e(target_e: int, count: int, start: date) -> list[date]:
    results = []
    d = start
    while len(results) < count and d.year <= start.year + 3:
        _, e, _ = day_info(d)
        if e == target_e:
            results.append(d)
        d += timedelta(days=1)
    return results


def next_dates_with_day_e_and_month_e(target_day_e: int, target_month_e: int, count: int, start: date) -> list[date]:
    results = []
    d = start
    while len(results) < count and d.year <= start.year + 3:
        _, e, b = day_info(d)
        if e == target_day_e and b["month"]["e"] == target_month_e:
            results.append(d)
        d += timedelta(days=1)
    return results


def pick_date_with_day_g(target_g: int, start: date) -> date:
    d = start
    while d.year <= start.year + 3:
        g, _, _ = day_info(d)
        if g == target_g:
            return d
        d += timedelta(days=1)
    raise RuntimeError("No date found for day.g")


def fmt_line(dt: date, g: int, e: int) -> str:
    return f"{dt.isoformat()} 08:00:00 (day.g={bz.gGodstem[g]}, day.e={bz.gEarthstem[e]})"


def main(out_path: str):
    lines: list[str] = []

    # Base years for date1 (all < 2025, spanning 1900–2020)
    base_years = [
        1905, 1912, 1928, 1937, 1944, 1956, 1967, 1977, 1988, 1998,
        2001, 2005, 2010, 2015, 2019,
    ]
    # Gaps (years) for date2 relative to date1 to simulate decade-apart cases
    gaps = [10, 12, 15, 18, 20, 25, 30, 14, 22, 28, 9, 16, 21, 24, 11]

    # Helper to get the first date on/after Jan 1 of a given year matching constraints
    def first_date_with_day_e_in_year_or_after(target_e: int, year_start: int) -> date:
        d = date(year_start, 1, 1)
        while True:
            _, e, _ = day_info(d)
            if e == target_e:
                return d
            d += timedelta(days=1)

    def first_date_with_day_e_and_month_e_in_year_or_after(target_day_e: int, target_month_e: int, year_start: int) -> date:
        d = date(year_start, 1, 1)
        while True:
            _, e, b = day_info(d)
            if e == target_day_e and b['month']['e'] == target_month_e:
                return d
            d += timedelta(days=1)

    def first_date_with_day_g_in_year_or_after(target_g: int, year_start: int) -> date:
        d = date(year_start, 1, 1)
        while True:
            g, _, _ = day_info(d)
            if g == target_g:
                return d
            d += timedelta(days=1)

    # 六合 pairs (day.e pairs)
    liuhe_pairs = [(5, 8), (0, 1), (3, 10), (4, 9), (6, 7), (2, 11), (5, 8), (0, 1), (3, 10), (4, 9)]
    lines.append("[liuhe]")
    for i in range(10):
        e1, e2 = liuhe_pairs[i]
        y1 = base_years[i % len(base_years)]
        y2 = y1 + gaps[i % len(gaps)]
        d1 = first_date_with_day_e_in_year_or_after(e1, y1)
        g1, e1a, _ = day_info(d1)
        d2 = first_date_with_day_e_in_year_or_after(e2, y2)
        g2, e2a, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, e1a)}, {d2.isoformat()} 11:00:00 (day.g={bz.gGodstem[g2]}, day.e={bz.gEarthstem[e2a]})")

    # 三合 pairs (ensure month.e forms the third member)
    triads = [(8, 0, 4), (11, 3, 7), (2, 6, 10), (5, 9, 1)]
    lines.append("\n[sanhe]")
    for i in range(10):
        a, b, c = triads[i % len(triads)]
        y1 = base_years[(i + 3) % len(base_years)]
        y2 = y1 + gaps[(i + 5) % len(gaps)]
        d1 = first_date_with_day_e_in_year_or_after(a, y1)
        g1, e1, _ = day_info(d1)
        d2 = first_date_with_day_e_and_month_e_in_year_or_after(b, c, y2)
        g2, e2, b2 = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, e1)}, {d2.isoformat()} 11:00:00 (day.g={bz.gGodstem[g2]}, day.e={bz.gEarthstem[e2]}, month.e={bz.gEarthstem[b2['month']['e']]})")

    # 天乙贵人 (stem-based). Use sg.MAP['tian_yi'] mapping day_g -> earth
    lines.append("\n[tian_yi]")
    ty_map = sg.MAP['tian_yi']
    # First: 阳贵人分支
    samples = [(g, ty_map[g][1]) for g in range(10)]  # pick 阳贵人 branch consistently
    for i, (g_needed, e_needed) in enumerate(samples):
        y1 = base_years[(i + 1) % len(base_years)]
        y2 = y1 + gaps[(i + 7) % len(gaps)]
        d1 = first_date_with_day_g_in_year_or_after(g_needed, y1)
        g1, e1, _ = day_info(d1)
        d2 = first_date_with_day_e_in_year_or_after(e_needed, y2)
        g2, e2, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, e1)}, {fmt_line(d2, g2, e2)}")
    # Then: 阴贵人分支（补齐另一个地支）
    samples_yin = [(g, ty_map[g][0]) for g in range(10)]
    for i, (g_needed, e_needed) in enumerate(samples_yin):
        y1 = base_years[(i + 9) % len(base_years)]
        y2 = y1 + gaps[(i + 3) % len(gaps)]
        d1 = first_date_with_day_g_in_year_or_after(g_needed, y1)
        g1, e1, _ = day_info(d1)
        d2 = first_date_with_day_e_in_year_or_after(e_needed, y2)
        g2, e2, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, e1)}, {fmt_line(d2, g2, e2)}")

    # 文昌贵人 (stem-based). sg.MAP['wenchang'] day_g -> earth
    lines.append("\n[wenchang]")
    wc_map = sg.MAP['wenchang']
    w_items = list(wc_map.items())[:10]
    for i, (g_needed, e_needed) in enumerate(w_items):
        y1 = base_years[(i + 5) % len(base_years)]
        y2 = y1 + gaps[(i + 9) % len(gaps)]
        d1 = first_date_with_day_g_in_year_or_after(g_needed, y1)
        g1, e1, _ = day_info(d1)
        d2 = first_date_with_day_e_in_year_or_after(e_needed, y2)
        g2, e2, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, e1)}, {fmt_line(d2, g2, e2)}")

    # 桃花 (branch-based). sg.MAP['taohua'] day_e -> earth
    lines.append("\n[taohua]")
    th_map = sg.MAP['taohua']
    th_items = list(th_map.items())[:10]
    for i, (e_user, e_date) in enumerate(th_items):
        y1 = base_years[(i + 2) % len(base_years)]
        y2 = y1 + gaps[(i + 4) % len(gaps)]
        d1 = first_date_with_day_e_in_year_or_after(e_user, y1)
        g1, e1, _ = day_info(d1)
        d2 = first_date_with_day_e_in_year_or_after(e_date, y2)
        g2, e2, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, e1)}, {fmt_line(d2, g2, e2)}")

    # Far apart examples (different years: 10/12/15/18/20/25/30 years etc.)
    def next_date_with_day_e_after_year(target_e: int, min_year: int) -> date:
        d = date(min_year, 1, 1)
        while True:
            _, e, _ = day_info(d)
            if e == target_e:
                return d
            d += timedelta(days=1)

    gaps = [10, 12, 15, 18, 20, 25, 30, 32, 35, 40]

    # Removed separate *_far sections; above lists already include wide year ranges

    # Edge cases: same date 22:59:59 vs 23:00:00 showing day pillar change affects relation
    lines.append("\n[edge_cases]")
    # Helper: compute day.e at two times on same civil date
    def day_e_at_times(d: date) -> tuple[int, int, int, int]:
        b1 = bz.getDateTimeGodEarthStem(d.year, d.month, d.day, 22, 59, 59)
        b2 = bz.getDateTimeGodEarthStem(d.year, d.month, d.day, 23, 0, 0)
        return b1['day']['g'], b1['day']['e'], b2['day']['g'], b2['day']['e']

    # Find a civil date where e@2259 != target and e@2300 == target
    def find_edge_date_for_target_e(target_e: int, start_year: int) -> date:
        d = date(start_year, 1, 1)
        limit = start_year + 5
        while d.year <= limit:
            _, e1, _, e2 = day_e_at_times(d)
            if e1 != target_e and e2 == target_e:
                return d
            d += timedelta(days=1)
        return d

    def fmt_edge_lines(title: str, date1: date, date2: date, form_at_2300: bool):
        # Print two lines: first (did not form at 22:59:59), second (form at 23:00:00)
        g1, e1, _, _ = day_e_at_times(date1)
        g2a, e2a, g2b, e2b = day_e_at_times(date2)
        lines.append(f"{title}-no:  {fmt_line(date1, g1, e1)}, {date2.isoformat()} 22:59:59 (day.g={bz.gGodstem[g2a]}, day.e={bz.gEarthstem[e2a]})")
        lines.append(f"{title}-yes: {fmt_line(date1, g1, e1)}, {date2.isoformat()} 23:00:00 (day.g={bz.gGodstem[g2b]}, day.e={bz.gEarthstem[e2b]})")

    def fmt_edge_lines_with_month(title: str, date1: date, date2: date):
        # Include month pillars at both 22:59:59 and 23:00:00 for date2
        g1, e1, _, _ = day_e_at_times(date1)
        b_a = bz.getDateTimeGodEarthStem(date2.year, date2.month, date2.day, 22, 59, 59)
        b_b = bz.getDateTimeGodEarthStem(date2.year, date2.month, date2.day, 23, 0, 0)
        g2a, e2a = b_a['day']['g'], b_a['day']['e']
        mg2a, me2a = b_a['month']['g'], b_a['month']['e']
        g2b, e2b = b_b['day']['g'], b_b['day']['e']
        mg2b, me2b = b_b['month']['g'], b_b['month']['e']
        lines.append(
            f"{title}-no:  {fmt_line(date1, g1, e1)}, {date2.isoformat()} 22:59:59 (day.g={bz.gGodstem[g2a]}, day.e={bz.gEarthstem[e2a]}, month.g={bz.gGodstem[mg2a]}, month.e={bz.gEarthstem[me2a]})"
        )
        lines.append(
            f"{title}-yes: {fmt_line(date1, g1, e1)}, {date2.isoformat()} 23:00:00 (day.g={bz.gGodstem[g2b]}, day.e={bz.gEarthstem[e2b]}, month.g={bz.gGodstem[mg2b]}, month.e={bz.gEarthstem[me2b]})"
        )

    # liuhe: choose date1 with day.e = 子(0) so target partner is 丑(1)
    date1_liuhe = first_date_with_day_e_in_year_or_after(0, 1988)
    edge_date_liuhe = find_edge_date_for_target_e(1, 1990)
    fmt_edge_lines("liuhe", date1_liuhe, edge_date_liuhe, form_at_2300=True)

    # sanhe: triad {申(8), 子(0), 辰(4)}; choose date1.e=申(8), date2 month.e=辰(4) and date2.e@2300=子(0)
    date1_sanhe = first_date_with_day_e_in_year_or_after(8, 1967)
    # Find a date2 where month.e=辰(4) and e@2300==子(0) and e@2259!=子(0)
    d = date(1991, 1, 1)
    while True:
        _, e1, _, e2 = day_e_at_times(d)
        b = bz.getCalendar10kGodEarthStem(d.year, d.month, d.day, 22, 59, 59)
        if b['month']['e'] == 4 and e2 == 0 and e1 != 0:
            date2_sanhe = d
            break
        d += timedelta(days=1)
    fmt_edge_lines_with_month("sanhe", date1_sanhe, date2_sanhe)

    # tian_yi: user day.g=甲(0) → target e=未(7)
    date1_tianyi = first_date_with_day_g_in_year_or_after(0, 1988)
    date2_tianyi = find_edge_date_for_target_e(7, 1992)
    fmt_edge_lines("tian_yi", date1_tianyi, date2_tianyi, form_at_2300=True)

    # wenchang: user day.g=甲(0) → target e=巳(5)
    date1_wenchang = first_date_with_day_g_in_year_or_after(0, 1988)
    date2_wenchang = find_edge_date_for_target_e(5, 1993)
    fmt_edge_lines("wenchang", date1_wenchang, date2_wenchang, form_at_2300=True)

    # taohua: user day.e=子(0) → target e=酉(9)
    date1_taohua = first_date_with_day_e_in_year_or_after(0, 1988)
    date2_taohua = find_edge_date_for_target_e(9, 1994)
    fmt_edge_lines("taohua", date1_taohua, date2_taohua, form_at_2300=True)

    # Negative examples (do NOT form the relation)
    # Helper: triads
    TRIADS = [ {8,0,4}, {11,3,7}, {2,6,10}, {5,9,1} ]

    lines.append("\n[liuhe_not]")
    # generate 10 pairs where e2 is not the liuhe partner of e1
    not_pairs = []
    for e1 in range(12):
        for e2 in range(12):
            if e1 == e2:
                continue
            if not bz.is6Harmony(e1, e2):
                not_pairs.append((e1, e2))
    for i in range(10):
        e1, e2_wrong = not_pairs[i * 5]
        y1 = base_years[i % len(base_years)]
        y2 = y1 + gaps[i % len(gaps)]
        d1 = first_date_with_day_e_in_year_or_after(e1, y1)
        g1, ee1, _ = day_info(d1)
        d2 = first_date_with_day_e_in_year_or_after(e2_wrong, y2)
        g2, ee2, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, ee1)}, {fmt_line(d2, g2, ee2)}")

    # 三合 not: choose e1 and e2 from different triads, thus not forming 半三合 or 三合 regardless of month
    lines.append("\n[sanhe_not]")
    triad_list = [set(t) for t in triads]
    diff_pairs = []
    for i_t, tset in enumerate(triad_list):
        other = triad_list[(i_t + 1) % len(triad_list)]
        for e1 in tset:
            for e2 in other:
                diff_pairs.append((e1, e2))
    for i in range(10):
        e1, e2 = diff_pairs[i * 3]
        y1 = base_years[(i + 1) % len(base_years)]
        y2 = y1 + gaps[(i + 2) % len(gaps)]
        d1 = first_date_with_day_e_in_year_or_after(e1, y1)
        g1, ee1, _ = day_info(d1)
        d2 = first_date_with_day_e_in_year_or_after(e2, y2)
        g2, ee2, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, ee1)}, {fmt_line(d2, g2, ee2)}")

    # 天乙贵人 not: pick day_g and e_date not equal to mapped ones
    lines.append("\n[tian_yi_not]")
    for i, day_g in enumerate(range(10)):
        y1 = base_years[(i + 2) % len(base_years)]
        y2 = y1 + gaps[(i + 6) % len(gaps)]
        d1 = first_date_with_day_g_in_year_or_after(day_g, y1)
        g1, e1, _ = day_info(d1)
        mapped = set(sg.MAP['tian_yi'][day_g].values())
        # pick an e not in mapped
        e_wrong = next(e for e in range(12) if e not in mapped)
        d2 = first_date_with_day_e_in_year_or_after(e_wrong, y2)
        g2, e2, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, e1)}, {fmt_line(d2, g2, e2)}")

    # 文昌 not: pick day_g and an e_date not equal to mapping
    lines.append("\n[wenchang_not]")
    w_items = list(sg.MAP['wenchang'].items())[:10]
    for i, (day_g, e_map) in enumerate(w_items):
        y1 = base_years[(i + 4) % len(base_years)]
        y2 = y1 + gaps[(i + 8) % len(gaps)]
        d1 = first_date_with_day_g_in_year_or_after(day_g, y1)
        g1, e1, _ = day_info(d1)
        e_wrong = next(e for e in range(12) if e != e_map)
        d2 = first_date_with_day_e_in_year_or_after(e_wrong, y2)
        g2, e2, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, e1)}, {fmt_line(d2, g2, e2)}")

    # 桃花 not: pick e_user and an e_date not equal to mapping
    lines.append("\n[taohua_not]")
    t_items = list(sg.MAP['taohua'].items())[:10]
    for i, (e_user, e_map) in enumerate(t_items):
        y1 = base_years[(i + 3) % len(base_years)]
        y2 = y1 + gaps[(i + 7) % len(gaps)]
        d1 = first_date_with_day_e_in_year_or_after(e_user, y1)
        g1, e1, _ = day_info(d1)
        e_wrong = next(e for e in range(12) if e != e_map)
        d2 = first_date_with_day_e_in_year_or_after(e_wrong, y2)
        g2, e2, _ = day_info(d2)
        lines.append(f"{fmt_line(d1, g1, e1)}, {fmt_line(d2, g2, e2)}")

    Path(out_path).write_text("\n".join(lines), encoding="utf-8")


if __name__ == "__main__":
    main("api/tests/bazi_relation_data.txt")


