"""Built-in pipeline templates for common bioinformatics analyses."""

# RNA-seq differential expression with DESeq2
RNASEQ_DESEQ2_TEMPLATE = """# RNA-seq Differential Expression Analysis Pipeline
# Generated by LLMomics v{llmomics_version}
# Date: {generated_date}

import os
from pathlib import Path
import pandas as pd

# Configuration
configfile: "config.yaml"

# Directories
FASTQ_DIR = Path(config.get("fastq_dir", "fastq"))
RESULTS_DIR = Path(config.get("results_dir", "results"))
QC_DIR = RESULTS_DIR / "qc"
ALIGN_DIR = RESULTS_DIR / "alignment"
COUNTS_DIR = RESULTS_DIR / "counts"
DESEQ_DIR = RESULTS_DIR / "deseq2"
LOG_DIR = Path(config.get("log_dir", "logs"))

# Create directories
for dir_path in [QC_DIR, ALIGN_DIR, COUNTS_DIR, DESEQ_DIR, LOG_DIR]:
    os.makedirs(dir_path, exist_ok=True)

# Load sample information
samples_df = pd.read_csv(config.get("samples", "samples.tsv"), sep="\\t")
SAMPLES = samples_df["sample_id"].tolist()

# Reference files
GENOME_FA = config.get("genome_fasta", "reference/genome.fa")
GENOME_GTF = config.get("genome_gtf", "reference/genes.gtf")
STAR_INDEX = config.get("star_index", "reference/star_index")

# Parameters
THREADS = config.get("threads", {{}})
DESEQ_PARAMS = config.get("deseq2_params", {{}})

# Rules
rule all:
    input:
        QC_DIR / "multiqc_report.html",
        DESEQ_DIR / "results_summary.html",
        DESEQ_DIR / "normalized_counts.csv",
        DESEQ_DIR / "degs_significant.csv"

rule fastqc:
    input:
        r1 = FASTQ_DIR / "{{sample}}_R1.fastq.gz",
        r2 = FASTQ_DIR / "{{sample}}_R2.fastq.gz"
    output:
        html1 = QC_DIR / "{{sample}}_R1_fastqc.html",
        html2 = QC_DIR / "{{sample}}_R2_fastqc.html",
        zip1 = QC_DIR / "{{sample}}_R1_fastqc.zip",
        zip2 = QC_DIR / "{{sample}}_R2_fastqc.zip"
    log:
        LOG_DIR / "fastqc_{{sample}}.log"
    threads: THREADS.get("fastqc", 2)
    conda:
        "envs/rnaseq.yaml"
    shell:
        "fastqc -t {{threads}} -o {{QC_DIR}} {{input.r1}} {{input.r2}} &> {{log}}"

rule trim_galore:
    input:
        r1 = FASTQ_DIR / "{{sample}}_R1.fastq.gz",
        r2 = FASTQ_DIR / "{{sample}}_R2.fastq.gz"
    output:
        r1 = temp(FASTQ_DIR / "trimmed" / "{{sample}}_R1_trimmed.fq.gz"),
        r2 = temp(FASTQ_DIR / "trimmed" / "{{sample}}_R2_trimmed.fq.gz"),
        report1 = QC_DIR / "{{sample}}_R1.fastq.gz_trimming_report.txt",
        report2 = QC_DIR / "{{sample}}_R2.fastq.gz_trimming_report.txt"
    log:
        LOG_DIR / "trim_galore_{{sample}}.log"
    threads: THREADS.get("trim", 4)
    conda:
        "envs/rnaseq.yaml"
    shell:
        '''
        trim_galore --paired --cores {{threads}} \\
            --output_dir {{FASTQ_DIR}}/trimmed \\
            --fastqc_args "--outdir {{QC_DIR}}" \\
            {{input.r1}} {{input.r2}} &> {{log}}
        '''

rule star_align:
    input:
        r1 = FASTQ_DIR / "trimmed" / "{{sample}}_R1_trimmed.fq.gz",
        r2 = FASTQ_DIR / "trimmed" / "{{sample}}_R2_trimmed.fq.gz",
        index = STAR_INDEX
    output:
        bam = ALIGN_DIR / "{{sample}}.Aligned.sortedByCoord.out.bam",
        log_final = ALIGN_DIR / "{{sample}}.Log.final.out"
    log:
        LOG_DIR / "star_{{sample}}.log"
    threads: THREADS.get("alignment", 8)
    conda:
        "envs/rnaseq.yaml"
    shell:
        '''
        STAR --runThreadN {{threads}} \\
            --genomeDir {{input.index}} \\
            --readFilesIn {{input.r1}} {{input.r2}} \\
            --readFilesCommand zcat \\
            --outFileNamePrefix {{ALIGN_DIR}}/{{wildcards.sample}}. \\
            --outSAMtype BAM SortedByCoordinate \\
            --outSAMunmapped Within \\
            --outSAMattributes Standard &> {{log}}
        '''

rule index_bam:
    input:
        ALIGN_DIR / "{{sample}}.Aligned.sortedByCoord.out.bam"
    output:
        ALIGN_DIR / "{{sample}}.Aligned.sortedByCoord.out.bam.bai"
    conda:
        "envs/rnaseq.yaml"
    shell:
        "samtools index {{input}}"

rule featurecounts:
    input:
        bams = expand(ALIGN_DIR / "{{sample}}.Aligned.sortedByCoord.out.bam", sample=SAMPLES),
        gtf = GENOME_GTF
    output:
        counts = COUNTS_DIR / "counts_matrix.txt",
        summary = COUNTS_DIR / "counts_matrix.txt.summary"
    log:
        LOG_DIR / "featurecounts.log"
    threads: THREADS.get("counting", 4)
    conda:
        "envs/rnaseq.yaml"
    shell:
        '''
        featureCounts -T {{threads}} \\
            -p -B -C \\
            -a {{input.gtf}} \\
            -o {{output.counts}} \\
            {{input.bams}} &> {{log}}
        '''

rule multiqc:
    input:
        expand(QC_DIR / "{{sample}}_R{{read}}_fastqc.zip", sample=SAMPLES, read=[1,2]),
        expand(ALIGN_DIR / "{{sample}}.Log.final.out", sample=SAMPLES),
        COUNTS_DIR / "counts_matrix.txt.summary"
    output:
        QC_DIR / "multiqc_report.html"
    log:
        LOG_DIR / "multiqc.log"
    conda:
        "envs/rnaseq.yaml"
    shell:
        "multiqc -f -o {{QC_DIR}} {{QC_DIR}} {{ALIGN_DIR}} {{COUNTS_DIR}} &> {{log}}"

rule deseq2:
    input:
        counts = COUNTS_DIR / "counts_matrix.txt",
        samples = config.get("samples", "samples.tsv")
    output:
        results = DESEQ_DIR / "deseq2_results.rds",
        norm_counts = DESEQ_DIR / "normalized_counts.csv",
        degs = DESEQ_DIR / "degs_all.csv",
        degs_sig = DESEQ_DIR / "degs_significant.csv",
        report = DESEQ_DIR / "results_summary.html"
    params:
        design = config.get("design_formula", "~ condition"),
        contrast = config.get("contrast", ["condition", "treatment", "control"]),
        padj_cutoff = DESEQ_PARAMS.get("padj_cutoff", 0.05),
        lfc_cutoff = DESEQ_PARAMS.get("lfc_cutoff", 1.0)
    log:
        LOG_DIR / "deseq2.log"
    conda:
        "envs/deseq2.yaml"
    script:
        "scripts/run_deseq2.R"
"""

# ChIP-seq peak calling with MACS2
CHIPSEQ_MACS2_TEMPLATE = """# ChIP-seq Peak Calling Pipeline
# Generated by LLMomics v{llmomics_version}
# Date: {generated_date}

import os
from pathlib import Path
import pandas as pd

# Configuration
configfile: "config.yaml"

# Directories setup
FASTQ_DIR = Path(config.get("fastq_dir", "fastq"))
RESULTS_DIR = Path(config.get("results_dir", "results"))
QC_DIR = RESULTS_DIR / "qc"
ALIGN_DIR = RESULTS_DIR / "alignment"
PEAKS_DIR = RESULTS_DIR / "peaks"
BIGWIG_DIR = RESULTS_DIR / "bigwig"
LOG_DIR = Path(config.get("log_dir", "logs"))

# Create directories
for dir_path in [QC_DIR, ALIGN_DIR, PEAKS_DIR, BIGWIG_DIR, LOG_DIR]:
    os.makedirs(dir_path, exist_ok=True)

# Load sample information
samples_df = pd.read_csv(config.get("samples", "samples.tsv"), sep="\\t")
SAMPLES = samples_df["sample_id"].tolist()
CONTROLS = dict(zip(samples_df["sample_id"], samples_df["control_id"]))

# Reference files
GENOME_FA = config.get("genome_fasta", "reference/genome.fa")
BOWTIE2_INDEX = config.get("bowtie2_index", "reference/bowtie2_index")

# Parameters
THREADS = config.get("threads", {{}})
MACS2_PARAMS = config.get("macs2_params", {{}})

rule all:
    input:
        QC_DIR / "multiqc_report.html",
        expand(PEAKS_DIR / "{{sample}}_peaks.narrowPeak", sample=SAMPLES),
        expand(BIGWIG_DIR / "{{sample}}.bw", sample=SAMPLES)

rule fastqc:
    input:
        FASTQ_DIR / "{{sample}}.fastq.gz"
    output:
        html = QC_DIR / "{{sample}}_fastqc.html",
        zip = QC_DIR / "{{sample}}_fastqc.zip"
    log:
        LOG_DIR / "fastqc_{{sample}}.log"
    threads: THREADS.get("fastqc", 2)
    conda:
        "envs/chipseq.yaml"
    shell:
        "fastqc -t {{threads}} -o {{QC_DIR}} {{input}} &> {{log}}"

rule bowtie2_align:
    input:
        fastq = FASTQ_DIR / "{{sample}}.fastq.gz",
        index = BOWTIE2_INDEX
    output:
        bam = ALIGN_DIR / "{{sample}}.sorted.bam"
    log:
        LOG_DIR / "bowtie2_{{sample}}.log"
    threads: THREADS.get("alignment", 8)
    conda:
        "envs/chipseq.yaml"
    shell:
        '''
        bowtie2 -p {{threads}} -x {{input.index}} -U {{input.fastq}} 2> {{log}} | \\
        samtools sort -@ {{threads}} -o {{output.bam}} -
        '''

rule index_bam:
    input:
        ALIGN_DIR / "{{sample}}.sorted.bam"
    output:
        ALIGN_DIR / "{{sample}}.sorted.bam.bai"
    conda:
        "envs/chipseq.yaml"
    shell:
        "samtools index {{input}}"

rule macs2_callpeak:
    input:
        treatment = ALIGN_DIR / "{{sample}}.sorted.bam",
        control = lambda wc: ALIGN_DIR / f"{{CONTROLS[wc.sample]}}.sorted.bam" if CONTROLS[wc.sample] else [],
        treatment_idx = ALIGN_DIR / "{{sample}}.sorted.bam.bai",
        control_idx = lambda wc: ALIGN_DIR / f"{{CONTROLS[wc.sample]}}.sorted.bam.bai" if CONTROLS[wc.sample] else []
    output:
        narrowPeak = PEAKS_DIR / "{{sample}}_peaks.narrowPeak",
        summits = PEAKS_DIR / "{{sample}}_summits.bed",
        xls = PEAKS_DIR / "{{sample}}_peaks.xls"
    params:
        genome_size = MACS2_PARAMS.get("genome_size", "hs"),
        qvalue = MACS2_PARAMS.get("qvalue", 0.05),
        control_arg = lambda wc: f"-c {{ALIGN_DIR}}/{{CONTROLS[wc.sample]}}.sorted.bam" if CONTROLS[wc.sample] else ""
    log:
        LOG_DIR / "macs2_{{sample}}.log"
    conda:
        "envs/chipseq.yaml"
    shell:
        '''
        macs2 callpeak -t {{input.treatment}} {{params.control_arg}} \\
            -f BAM -g {{params.genome_size}} \\
            -n {{wildcards.sample}} \\
            --outdir {{PEAKS_DIR}} \\
            -q {{params.qvalue}} &> {{log}}
        '''

rule bamCoverage:
    input:
        bam = ALIGN_DIR / "{{sample}}.sorted.bam",
        bai = ALIGN_DIR / "{{sample}}.sorted.bam.bai"
    output:
        BIGWIG_DIR / "{{sample}}.bw"
    log:
        LOG_DIR / "bamCoverage_{{sample}}.log"
    threads: THREADS.get("bigwig", 4)
    conda:
        "envs/chipseq.yaml"
    shell:
        '''
        bamCoverage -b {{input.bam}} -o {{output}} \\
            --binSize 10 --normalizeUsing RPKM \\
            -p {{threads}} &> {{log}}
        '''

rule multiqc:
    input:
        expand(QC_DIR / "{{sample}}_fastqc.zip", sample=SAMPLES),
        expand(PEAKS_DIR / "{{sample}}_peaks.xls", sample=SAMPLES)
    output:
        QC_DIR / "multiqc_report.html"
    log:
        LOG_DIR / "multiqc.log"
    conda:
        "envs/chipseq.yaml"
    shell:
        "multiqc -f -o {{QC_DIR}} {{QC_DIR}} {{ALIGN_DIR}} {{PEAKS_DIR}} &> {{log}}"
"""

# Template dictionary
PIPELINE_TEMPLATES = {
    "rnaseq_deseq2": RNASEQ_DESEQ2_TEMPLATE,
    "chipseq_macs2": CHIPSEQ_MACS2_TEMPLATE,
}


# Pipeline template base class
class PipelineTemplate:
    """Base class for pipeline templates."""

    def __init__(self, name: str, template: str):
        self.name = name
        self.template = template

    def render(self, **context) -> str:
        """Render the template with given context."""
        return self.template.format(**context)
