#!/usr/bin/env python3
import glob
import pandas as pd


def load_snps(filtered_folders):
    """Load snps.tab files only from folders containing PT1 or PT4."""
    snps_files = [
        f for f in glob.glob("**/snps.tab", recursive=True)
        if any(tag in f for tag in filtered_folders)
    ]

    if len(snps_files) == 0:
        print(f"No snps.tab files found for {filtered_folders}")
        return None, None

    all_snps = []

    for f in snps_files:
        df = pd.read_csv(f, sep="\t", dtype=str)
        df["source"] = f

        # NEW: SNP identity only by genomic position (CHROM + POS)
        df["SNP_key"] = df["CHROM"] + "_" + df["POS"]

        all_snps.append(df)

    combined = pd.concat(all_snps, ignore_index=True)
    total = len(snps_files)
    return combined, total



def summarize(df, total, output_file):
    """Summarize SNPs by position, including allele diversity."""

    # Base count: how many genomes contain ANY SNP at this position
    base = (
        df.groupby("SNP_key")
          .agg(
              count=("source", "nunique"),
              chrom=("CHROM", "first"),
              pos=("POS", "first")
          )
    )

    # NEW: Summarize REF>ALT allele variants and their counts
    def allele_summary(group):
        alleles = []
        unique = group.drop_duplicates(subset=["REF", "ALT"])
        for _, row in unique.iterrows():
            alt = row["ALT"]
            ref = row["REF"]
            n = (group["ALT"] == alt).sum()
            alleles.append(f"{ref}>{alt}({n})")
        return ";".join(alleles)

    allele_info = df.groupby("SNP_key").apply(allele_summary)
    allele_info = allele_info.rename("alleles")

    base["percentage"] = (base["count"] / total) * 100

    # Merge allele diversity information
    result = base.join(allele_info)

    # Sort by how common the SNP is in the group
    result = result.sort_values(by="count", ascending=False)

    result.to_csv(output_file, sep="\t")
    return result



print("\n=== Loading PT4 SNPs ===")
PT4_df, PT4_total = load_snps(["PT4"])

print("\n=== Loading PT1 SNPs ===")
PT1_df, PT1_total = load_snps(["PT1"])

if PT4_df is None or PT1_df is None:
    exit()



# Generate per-group summaries
PT4_counts = summarize(PT4_df, PT4_total, "PT4_summary.tsv")
PT1_counts = summarize(PT1_df, PT1_total, "PT1_summary.tsv")



# Define SNP_key sets
PT4_keys = set(PT4_counts.index)
PT1_keys = set(PT1_counts.index)

shared = PT4_keys & PT1_keys
unique_PT4 = PT4_keys - PT1_keys
unique_PT1 = PT1_keys - PT4_keys



# Export unique SNPs
PT4_counts.loc[list(unique_PT4)].to_csv("unique_PT4_snps.tsv", sep="\t")
PT1_counts.loc[list(unique_PT1)].to_csv("unique_PT1_snps.tsv", sep="\t")



# Export shared SNPs with side-by-side counts
shared_df = (
    PT4_counts.loc[list(shared)]
        .join(PT1_counts.loc[list(shared)], lsuffix="_PT4", rsuffix="_PT1")
)

shared_df.to_csv("shared_snps.tsv", sep="\t")



print("\n=== DONE ===")
print(f"PT4 genomes analyzed: {PT4_total}")
print(f"PT1 genomes analyzed: {PT1_total}")
print("Generated files:")
print("  - PT4_summary.tsv")
print("  - PT1_summary.tsv")
print("  - shared_snps.tsv")
print("  - unique_PT4_snps.tsv")
print("  - unique_PT1_snps.tsv")