#!/usr/bin/env python3
import gzip
import re
from bisect import bisect_left
from collections import defaultdict
import numpy as np
import matplotlib.pyplot as plt

# -----------------------
# PATHS (your real files)
# -----------------------

TE_BED = "/BLUES/eric/ONT/hg38_TE_noY.bed.gz"

MC_FILES = {
    "Primed": "/BLUES/eric/ONT_WGBS/TSS_TES_plot/bedgraph/W3_primed.CG.cov10.bg",
    "Naive":  "/BLUES/eric/ONT_WGBS/TSS_TES_plot/bedgraph/W3_naive.CG.cov10.bg",
    "TSC":    "/BLUES/eric/ONT_WGBS/TSS_TES_plot/bedgraph/W3_TSC.CG.cov10.bg"
}

HMC_FILES = {
    "Primed": "/BLUES/eric/ONT/5hmC_bg/Primed_W3MECP2_aln.5hmC.cov10.bg",
    "Naive":  "/BLUES/eric/ONT/5hmC_bg/091724_Naive_W3MECPC2_aln.5hmC.cov10.bg",
    "TSC":    "/BLUES/eric/ONT/5hmC_bg/TSC_GFP_W3MECP2_merged_aln.5hmC.cov10.bg"
}

OUTFILE = "/BLUES/eric/ONT_WGBS/Heatmap/Heatmap_analysis/L1HS_metagene_Phase1_pm2kb_cov10.pdf"


# -----------------------
# PARAMETERS
# -----------------------

L1HS_LEN = 6064
MIN_FRAC = 0.8
MIN_LEN = int(L1HS_LEN * MIN_FRAC)

UP = 2000
DOWN = 2000

NBIN_UP = 40
NBIN_BODY = 120
NBIN_DOWN = 40

TOTAL = NBIN_UP + NBIN_BODY + NBIN_DOWN

SMOOTH = 5


# -----------------------
# TE parsing
# -----------------------

def parse_subfam(info):
    m = re.search(r'repeat_id\s+"([^"]+)"', info)
    if not m:
        return None
    return m.group(1).split(",")[0]


def load_l1hs():
    loci = []

    with gzip.open(TE_BED, "rt") as f:
        for line in f:
            chrn, s, e, strand, info = line.strip().split("\t")[:5]

            if parse_subfam(info) != "L1HS":
                continue

            s = int(s)
            e = int(e)

            if e - s < MIN_LEN:
                continue

            if strand not in "+-":
                continue

            loci.append((chrn, s, e, strand))

    print("L1HS loci:", len(loci))
    return loci


# -----------------------
# bedGraph loader
# -----------------------

def load_bg(path):
    pos = defaultdict(list)
    val = defaultdict(list)

    with open(path) as f:
        for line in f:
            c, s, e, m = line.strip().split("\t")[:4]

            s = int(s)
            m = float(m)

            pos[c].append(s)
            val[c].append(m)

    return pos, val


# -----------------------
# mean in interval
# -----------------------

def mean_interval(p, v, a, b):
    i0 = bisect_left(p, a)
    i1 = bisect_left(p, b)

    if i1 <= i0:
        return None

    s = 0
    n = 0

    for i in range(i0, i1):
        s += v[i]
        n += 1

    return s / n if n else None


# -----------------------
# profile computation
# -----------------------

def compute_profile(loci, pos, val):

    mD = [0.0] * TOTAL
    cnt = [0] * TOTAL

    BIN_UP = UP / NBIN_UP
    BIN_DOWN = DOWN / NBIN_DOWN

    for chrn, start, end, strand in loci:

        if chrn not in pos:
            continue

        P = pos[chrn]
        V = val[chrn]

        body_len = end - start
        k = body_len / NBIN_BODY

        if strand == "+":

            for i in range(NBIN_UP):
                a = start - UP + i * BIN_UP
                b = a + BIN_UP

                mv = mean_interval(P, V, int(a), int(b))
                if mv is None:
                    continue

                mD[i] += mv
                cnt[i] += 1

            for i in range(NBIN_BODY):
                a = start + i * k
                b = start + (i + 1) * k

                mv = mean_interval(P, V, int(a), int(b))
                if mv is None:
                    continue

                idx = NBIN_UP + i
                mD[idx] += mv
                cnt[idx] += 1

            for i in range(NBIN_DOWN):
                a = end + i * BIN_DOWN
                b = a + BIN_DOWN

                mv = mean_interval(P, V, int(a), int(b))
                if mv is None:
                    continue

                idx = NBIN_UP + NBIN_BODY + i
                mD[idx] += mv
                cnt[idx] += 1

        else:

            for i in range(NBIN_UP):
                a = end + i * BIN_UP
                b = a + BIN_UP

                mv = mean_interval(P, V, int(a), int(b))
                if mv is None:
                    continue

                mD[i] += mv
                cnt[i] += 1

            for i in range(NBIN_BODY):
                a = end - (i + 1) * k
                b = end - i * k

                mv = mean_interval(P, V, int(a), int(b))
                if mv is None:
                    continue

                idx = NBIN_UP + i
                mD[idx] += mv
                cnt[idx] += 1

            for i in range(NBIN_DOWN):
                a = start - DOWN + i * BIN_DOWN
                b = a + BIN_DOWN

                mv = mean_interval(P, V, int(a), int(b))
                if mv is None:
                    continue

                idx = NBIN_UP + NBIN_BODY + i
                mD[idx] += mv
                cnt[idx] += 1

    out = []

    for i in range(TOTAL):
        out.append(mD[i] / cnt[i] if cnt[i] else np.nan)

    return out


# -----------------------
# smoothing
# -----------------------

def smooth(v):

    if SMOOTH <= 1:
        return v

    w = np.ones(SMOOTH) / SMOOTH
    return np.convolve(v, w, mode="same")


# -----------------------
# main
# -----------------------

loci = load_l1hs()

profiles_mc = {}
profiles_hmc = {}

for cond in MC_FILES:

    print("Processing 5mC:", cond)

    pos, val = load_bg(MC_FILES[cond])
    p = compute_profile(loci, pos, val)

    profiles_mc[cond] = smooth(p)

for cond in HMC_FILES:

    print("Processing 5hmC:", cond)

    pos, val = load_bg(HMC_FILES[cond])
    p = compute_profile(loci, pos, val)

    profiles_hmc[cond] = smooth(p)


# -----------------------
# plot
# -----------------------

x = list(range(TOTAL))

colors = {
    "Primed": "#E64B35",
    "Naive": "#4DBBD5",
    "TSC": "#00A087"
}

fig, ax1 = plt.subplots(figsize=(12,5))
ax2 = ax1.twinx()

for cond in profiles_mc:
    ax1.plot(x, profiles_mc[cond], color=colors[cond], linewidth=2, label=cond + " 5mC")

for cond in profiles_hmc:
    ax2.plot(x, profiles_hmc[cond], linestyle="--", color=colors[cond], linewidth=2, label=cond + " 5hmC")

ax1.axvline(NBIN_UP, color="gray")
ax1.axvline(NBIN_UP + NBIN_BODY, color="gray")

ax1.set_xticks([0, NBIN_UP, NBIN_UP+NBIN_BODY, TOTAL-1])
ax1.set_xticklabels(["-2kb","5'","3'","+2kb"])

ax1.set_ylabel("5mC", fontsize=16)
ax2.set_ylabel("5hmC", fontsize=16)

ax1.set_title("L1HS metagene methylation (±2kb, cov≥10)", fontsize=18)

h1,l1 = ax1.get_legend_handles_labels()
h2,l2 = ax2.get_legend_handles_labels()

ax1.legend(h1+h2,l1+l2,frameon=False,bbox_to_anchor=(1.02,0.5),loc="center left")

plt.tight_layout()
plt.savefig(OUTFILE)
plt.close()

print("DONE:", OUTFILE)