#!/usr/bin/env nextflow


process CONPLEX {
    //debug true

    publishDir "${params.outdir}/${params.project_name}/2_conplex", mode: 'copy'
    
    memory { params.conplex_initial_memory.toInteger().GB + (task.attempt - 1) * params.conplex_growth_memory.toInteger().GB }
    errorStrategy { task.attempt <= params.conplex_max_retries.toInteger() ? 'retry' : params.conplex_fail_action }
    maxRetries params.conplex_max_retries.toInteger()
    
    container "${params.container_mass_screen}"
    // containerOptions "--rm --gpus all"
    // maxForks 1    
    // errorStrategy 'ignore'
    
    afterScript "if [[ ${params.keep_enst} == 'false' ]]; then find . -name '*_results.zarr' -type d -exec rm -rf {} + 2>/dev/null || true; fi"
    
    input:
     tuple(
        val( id ),
        path( input_drug, arity: '1') ,    // CSV file with SMILES column
        path( input_metabol, arity: '1') // TSV file with SMILES and metabolite IDs
      )
     path input_zarr, arity: '1..*'    // ZARR DB holding all the protein vectors and their names
    
    output:
     tuple(
        val ( id ),
        path( "*drug_scores.tsv"),
        path( "*significant_interactions.tsv")
     )
    
    script:
    """
    set -e
    
    # Set up trap to clean up on any exit (success, failure, or signal)
    if [[ ${params.keep_enst} == 'false' ]]; then
        trap 'find \$PWD -name "*_results.zarr" -type d -exec rm -rf {} + 2>/dev/null || true' EXIT
    fi
    
    #remove unnecessary data to run conplex 
    head -n 2 ${input_drug} > run.csv

    # Convert drug CSV and metabolite TSV to smiles.smi format
    python /app/convert.py run.csv ${input_metabol} -o smiles.smi
    
    # Run screening for each zarr database
    for zarr_dir in ${input_zarr}; do
        python /app/screen.py -i smiles.smi -z "\${zarr_dir}"
    done
    
    # Process results and generate final outputs
    python /app/get_round_2.py --threshold ${params.threshold} --workdir \$PWD --round 2 --drug-csv ${input_drug}
    """
}

process MERGE_DRUG {
    //debug true
    
    memory '10 GB'
    
    container "${params.container_mass_screen}"
    // maxForks 1    
    // errorStrategy 'ignore'
    
    input:
     path 'inputs/drug_*.tsv', arity: '1..*'
    
    output:
     path "patient_0_drug_scores.tsv"
    
    script:
"""
#!/opt/conda/bin/python
import pandas as pd
from pathlib import Path

inputs_dir = Path("inputs")
tsv_files = sorted(inputs_dir.glob("drug_*.tsv"))

dfs = [pd.read_csv(f, sep="\\t") for f in tsv_files]
df_out = pd.concat(dfs, ignore_index=True)

df_out.to_csv("patient_0_drug_scores.tsv", sep="\\t", index=False)
"""
}

process MERGE_INTERACTIONS {
    //debug true
    
    memory '10 GB'
    
    container "${params.container_mass_screen}"
    // maxForks 1    
    // errorStrategy 'ignore'
    
    input:
     path 'inputs/interaction_*.tsv'
    
    output:
     path "patient_0_significant_interactions.tsv"
    
    script:
"""
#!/opt/conda/bin/python
import pandas as pd
from pathlib import Path

inputs_dir = Path("inputs")
tsv_files = sorted(inputs_dir.glob("interaction_*.tsv"))

dfs = [pd.read_csv(f, sep="\\t") for f in tsv_files]
df_out = pd.concat(dfs, ignore_index=True)

df_out.to_csv("patient_0_significant_interactions.tsv", sep="\\t", index=False)
"""
}


process NETWORK_ENRICHMENT {
  memory { params.string_initial_memory.toInteger().GB + (task.attempt - 1) * params.string_growth_memory.toInteger().GB }

  container "${params.container_conplex}"
  // containerOptions "${params.containerOptions}"
  publishDir "${params.outdir}/${params.project_name}/3_string", mode: 'copy'
  // // Temporarily disabled debug prints
  // debug true
  maxForks params.string_max_forks.toInteger() ?: null
  errorStrategy { task.attempt <= params.string_max_retries.toInteger() ? 'retry' : params.string_fail_action }
  maxRetries params.string_max_retries.toInteger()

  input:
    path interactions

  output:
    path "*_network_enrichment.tsv", emit: network_enrichment
    path "*_network_interactions.tsv", emit: network_interactions

  script:
    """
#!/opt/conda/envs/conplex-dti/bin/python
import sys
import pandas as pd
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

# Read ONLY the columns we use. The 'smile' column in significant_interactions.tsv
# is the bulk of the file (long SMILES strings × millions of rows for metabolite-heavy
# drugs → ~2 GB on disk), and we never reference it here. Dropping it on read cuts
# the in-memory DataFrame from ~2 GB to ~350 MB on worst-case drugs and prevents OOM.
interaction_data = pd.read_csv('$interactions', sep='\\t', usecols=['transcipt', 'conplex_score'])
interaction_data = interaction_data.sort_values('conplex_score', ascending=False)

# string-db accepts at most 450 identifiers per call.
interaction_data = interaction_data.drop_duplicates('transcipt')
if interaction_data.shape[0] > 450:
    interaction_data = interaction_data.iloc[:450]

proteins = list(pd.unique(
    interaction_data[interaction_data['conplex_score'] > ${params.protein_network_threshold}]['transcipt']
))
output_name = '$interactions'.split('/')[-1].split('_significant_interactions.tsv')[0]

# Empty-identifier early exit: nothing to enrich; write empty outputs so the
# task succeeds without hitting string-db (which would otherwise return an
# error for an empty identifier list and trip a Nextflow retry).
if not proteins:
    print(f"NETWORK_ENRICHMENT: no proteins above threshold for {output_name}", flush=True)
    pd.DataFrame().to_csv(f'{output_name}_network_enrichment.tsv', sep='\\t')
    pd.DataFrame().to_csv(f'{output_name}_network_interactions.tsv', sep='\\t')
    sys.exit(0)

# urllib3 Retry handles 429 / 5xx / connection blips transparently at the
# adapter layer — no exception reaches user code, and Nextflow does not see a
# task failure. backoff_factor=2 produces delays 2,4,8,16,32 s (capped at
# Retry.BACKOFF_MAX), respecting any Retry-After header string-db sends.
# This absorbs the common rate-limit case WITHOUT spawning a fresh workdir,
# which would just get throttled the same way.
session = requests.Session()
session.mount('https://', HTTPAdapter(max_retries=Retry(
    total=5,
    backoff_factor=2,
    status_forcelist=[429, 500, 502, 503, 504],
)))


def fetch(url, kind):
    r = session.get(url, timeout=60)
    r.raise_for_status()
    if not r.text.strip():
        raise RuntimeError(f"{kind}: 200 but empty body from string-db")
    lines = r.text.split('\\n')
    data = [l.split('\\t') for l in lines]
    if not data or not data[0]:
        raise RuntimeError(f"{kind}: malformed response (no header row)")
    return pd.DataFrame(data[1:-1], columns=data[0])


protein_network_list = '%0d'.join(proteins)
url_network = f'https://string-db.org/api/tsv/network?identifiers={protein_network_list}&species=9606'
url_enrichment = f'https://string-db.org/api/tsv/enrichment?identifiers={protein_network_list}&species=9606'

enrichment_df = fetch(url_enrichment, 'enrichment')
interactions_df = fetch(url_network, 'network')

enrichment_df.to_csv(f'{output_name}_network_enrichment.tsv', sep='\\t')
interactions_df.to_csv(f'{output_name}_network_interactions.tsv', sep='\\t')
"""
}


process PREPROCESS_PROTEIN {
    
    memory '10 GB'
    accelerator 1
    
    container "${params.container_preprocess}"
    label 'gpu_process'
    // containerOptions "--rm --gpus all"

    // errorStrategy 'ignore'
    
    input:
      path("input_*.fasta", arity: '1..*')
    
    output:
      path "mutated.zarr"
    
    script:
    """
    set -e
    cat input_*.fasta > merged.fasta
    python /app/project.py merged.fasta mutated.zarr
    """
}