import os
import random

def generate_random_color():
    """Generate a random UCSC-compatible RGB color."""
    return f"{random.randint(50, 255)},{random.randint(50, 255)},{random.randint(50, 255)}"

def get_replicate_group(filename):
    """Extract the replicate group from a filename (heuristic-based)."""
    parts = filename.replace(".bw", "").replace("-", "_").split("_")
    if "rep" in parts:
        idx = parts.index("rep")
        return "_".join(parts[:idx+2])  # Include 'rep' and the number following it
    return "_".join(parts[:-1]) if len(parts) > 1 else parts[0]

def create_trackhub():
    # Ask for user inputs
    hub_name = input("Enter the hub name: ").strip()
    genome = input("Enter the genome name (e.g., hg19, hg38): ").strip()
    experiment_type = input("Enter the experiment type (e.g., iCLIP, ChIP-seq, RNA-seq): ").strip()

    # Define directory structure
    base_dir = os.getcwd()  # Current directory
    genome_dir = os.path.join(base_dir, genome)
    
    # Ensure genome directory exists
    os.makedirs(genome_dir, exist_ok=True)

    # Create hub.txt
    hub_txt_path = os.path.join(base_dir, "hub.txt")
    with open(hub_txt_path, "w") as f:
        f.write(f"""hub {hub_name}
shortLabel {hub_name}
longLabel {hub_name} - {experiment_type} Hub
genomesFile genomes.txt
email your_email@example.com
""")

    # Create genomes.txt
    genomes_txt_path = os.path.join(base_dir, "genomes.txt")
    with open(genomes_txt_path, "w") as f:
        f.write(f"""genome {genome}
trackDb {genome}/trackDb.txt
""")

    # Create trackDb.txt
    trackDb_txt_path = os.path.join(genome_dir, "trackDb.txt")
    bigwig_files = [f for f in os.listdir(genome_dir) if f.endswith((".bw", ".bigWig"))]

    # Assign colors to replicate groups
    replicate_colors = {}
    with open(trackDb_txt_path, "w") as f:
        for bw in bigwig_files:
            track_name = bw.replace("-", "_").replace(".bw", "")
            replicate_group = get_replicate_group(bw)

            if replicate_group not in replicate_colors:
                replicate_colors[replicate_group] = generate_random_color()
            
            color = replicate_colors[replicate_group]

            f.write(f"""track {track_name}
bigDataUrl {bw}
shortLabel {track_name}
longLabel {track_name} ({experiment_type})
type bigWig
visibility full
priority 1
autoScale on
windowingFunction mean
smoothingWindow 3
color {color}

""")

    print(f"✅ Track hub '{hub_name}' for {experiment_type} created successfully in {base_dir}.")
    print(f"📂 Structure:\n- hub.txt\n- genomes.txt\n- {genome}/trackDb.txt\n- {genome}/(bigWig files)")

if __name__ == "__main__":
    create_trackhub()