import os
CONFIG_YAML = """\
samples:
  sample1: 
    R1: "path to sample1 R1" 
    R2: "path to sample1 R2" 
  sample2: 
    R1: "path to sample2 R1" 
    R2: "path to sample2 R2" 
reference: "./Reference/mm10/genome"  
gtf: "./Reference/gencode.vM25.primary_assembly.annotation.gtf.gz"  
counts: "./Counts"  
threads: 8   
"""
SNAKEFILE = """\
configfile: "config.yaml"
rule all:
    input:
        "QCresult/multiqc_report.html",  
        expand("Cleandata/{sample}_1_val_1.fq", sample=config["samples"]),
        expand("Cleandata/{sample}_2_val_2.fq", sample=config["samples"]),
        "Mapping/multiqc_report.html",  
        "Counts/counts.txt",             
        "Cleandata/multiqc_report.html"  
rule fastqc:
    input:
        fq="Rawdata/{fq}"
    output:
        zip="QCresult/{wildcards.fq}.fastqc.zip",
        html="QCresult/{wildcards.fq}.fastqc.html"
        report="QCresult/multiqc_report.html"
    threads: config["threads"]
    shell:
        "fastqc -t {threads} -o QCresult {input.fq}"
        "multiqc -o QCresult QCresult/*.zip"

rule trim_galore:
    input:
        r1=lambda wildcards: config["samples"][wildcards.sample]["R1"],
        r2=lambda wildcards: config["samples"][wildcards.sample]["R2"]
    output:
        r1_clean="Cleandata/{sample}_1_val_1.fq",
        r2_clean="Cleandata/{sample}_2_val_2.fq"
        report="Cleandata/multiqc_report.html"
    threads: config["threads"]
    shell:
        "trim_galore -q 20 --length 36 --max_n 3 --stringency 3 --fastqc --paired -o Cleandata {input.r1} {input.r2}"
        "multiqc -o Cleandata Cleandata/*.zip"
rule hisat2:
    input:
        r1_clean="Cleandata/{sample}_1_val_1.fq",
        r2_clean="Cleandata/{sample}_2_val_2.fq"
    output:
        bam="Mapping/{sample}.Hisat_aln.sorted.bam",
        bai="Mapping/{sample}.Hisat_aln.sorted.bam.bai"
        report="Mapping/multiqc_report.html"
    threads: config["threads"]
    shell:
        "hisat2 -p {threads} -x {config[reference]} -1 {input.r1_clean} -2 {input.r2_clean} 2> Mapping/{wildcards.sample}.log | samtools sort -@ {threads} -o {output.bam}"
        "samtools index {output.bam} {output.bai}"
        "multiqc -o Mapping Mapping/*.log"

rule featureCounts:
    input:
        bams=expand("Mapping/{sample}.Hisat_aln.sorted.bam", sample=config["samples"]),
        gtf=config["gtf"]
    output:
        counts="Counts/all.id.txt",
        summary="Counts/all.id.txt.summary"
    threads: config["threads"]
    shell:
        "featureCounts -T {threads} -p -t exon -g gene_id -a {input.gtf} -o {output.counts} {input.bams}"
rule counts_matrix:
    input:
        "Counts/all.id.txt"
    output:
        "Counts/counts.txt"
    shell:
        "cat {input} | cut -f1,7- > {output}"
"""
def init_project(base_dir):
    """
    Create necessary files and folders for the RNA-seq pipeline.
    """
    base_dir = os.path.abspath(base_dir)
    folders = ["Rawdata", "Cleandata", "QCresult", "Mapping", "Counts", "Reference"]
    for folder in folders:
        os.makedirs(os.path.join(base_dir, folder), exist_ok=True)
    with open(os.path.join(base_dir, "config.yaml"), "w") as f:
        f.write(CONFIG_YAML)
    with open(os.path.join(base_dir, "Snakefile"), "w") as f:
        f.write(SNAKEFILE)
    print(f"Pipeline initialized in: {base_dir}")

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Generate Snakemake project structure for RNA-seq.")
    parser.add_argument("base_dir", help="The root directory where the pipeline will be initialized.")
    args = parser.parse_args()
    init_project(args.base_dir)


