# -*- coding: utf-8 -*-
"""SELECT-NET_.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1sv9fiafPeB_FdGLnBNzF1ZYJtxliU-jE
"""

#                       Online Resource 6
# ==============================================================================
# SELECT-NET
# Google Colab pipeline
#Click on play buttom tu run the pipeline - it will request at the bottom ...
# ... the SRA number and then it wil Run the FULL mode for LITE mode see the
# Supplementary SOP material.
#
# - Drive Checkpoints
# - RPKM abundance
# - Trim Galore + Cutadapt
# - Prokka with fallback to Prodigal
# - KOfam + KEGG pathway mapping
# ==============================================================================

import os
import re
import sys
import time
import glob
import json
import shutil
import subprocess
from pathlib import Path
import requests
import glob
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt

# ==============================================================================
# 0A) CLONAR REPO DE SELECT-NET DESDE GITHUB (para referencias externas core)
# ==============================================================================

import os
import subprocess

GITHUB_REPO_URL = "https://github.com/amoresgr/SELECT-NET.git"
LOCAL_REPO_DIR = "/content/SELECT-NET"
CORE_REF_DIR = os.path.join(LOCAL_REPO_DIR, "Core_EnteroDB")

print("\n loading core-External references from GitHub...")

if not os.path.exists(LOCAL_REPO_DIR):
    print("\n loading repository at SELECT-NET...")
    subprocess.run(["git", "clone", GITHUB_REPO_URL, LOCAL_REPO_DIR], check=True)
else:
    print("\n repository already exists in Colab. Updating...")
    subprocess.run(["git", "-C", LOCAL_REPO_DIR, "pull"], check=False)

print(f"Repository Ready: {LOCAL_REPO_DIR}")
print(f"Reference Core Folder: {CORE_REF_DIR}")

# ==============================================================================
# 0) CONFIGURACIÓN
# ==============================================================================

REQUIRE_PROKKA = True   # True = si Prokka falla, el pipeline se detiene
FORCE_RERUN = False             # True = ignora checkpoints y rehace todo
SAVE_TRIM_CHECKPOINTS = False   # True = guarda clean FASTQ en Drive (ocupa mucho espacio)
SAVE_BAM_CHECKPOINTS = False    # True = guarda BAM indexado en Drive (ocupa mucho espacio)
CLEANUP_AFTER_SUCCESS = False   # False mientras estabilizamos el pipeline

# ==============================================================================
# 1) UTILIDADES
# ==============================================================================

def run(cmd, check=True, capture_output=False, text=True, shell=False):
    print(f"\n[CMD] {cmd if isinstance(cmd, str) else ' '.join(cmd)}")
    return subprocess.run(
        cmd,
        check=check,
        capture_output=capture_output,
        text=text,
        shell=shell
    )

def ensure_dir(path):
    os.makedirs(path, exist_ok=True)

def first_existing(paths):
    for p in paths:
        if p and os.path.exists(p):
            return p
    return None

def parse_gff_attributes(attr_str):
    d = {}
    for item in attr_str.split(';'):
        item = item.strip()
        if not item:
            continue
        if '=' in item:
            k, v = item.split('=', 1)
            d[k.strip()] = v.strip()
    return d

def cleanup_large_files(files_to_remove):
    for f in files_to_remove:
        try:
            if f and os.path.exists(f):
                if os.path.isdir(f):
                    shutil.rmtree(f)
                else:
                    os.remove(f)
                print(f"Eliminado: {f}")
        except Exception as e:
            print(f"No se pudo eliminar {f}: {e}")

def sanitize_fasta_headers(in_fasta, out_fasta):
    with open(in_fasta, 'r') as fin, open(out_fasta, 'w') as fout:
        c = 1
        for line in fin:
            if line.startswith('>'):
                fout.write(f">contig_{c}\n")
                c += 1
            else:
                fout.write(line)

def check_prokka_preflight(prokka_bin):
    print("\nPreflight de Prokka...")

    if not prokka_bin or not os.path.exists(prokka_bin):
        print("No hay ejecutable Prokka disponible.")
        return False

    for cmd in [[prokka_bin, '--version'], [prokka_bin, '--depends'], [prokka_bin, '--listdb']]:
        try:
            res = subprocess.run(cmd, capture_output=True, text=True, check=False)
            print(f"\n[PROKKA CHECK] {' '.join(cmd)}")
            print(res.stdout if res.stdout else res.stderr)
        except Exception as e:
            print(f"No se pudo ejecutar {' '.join(cmd)}: {e}")

    return True

def parse_annotation_gff(gff_path, gene_attr):
    """
    Parsea GFF de Prokka o Prodigal y devuelve metadatos por CDS.
    gene_attr suele ser:
      - 'locus_tag' (Prokka)
      - 'ID'        (Prodigal)
    """
    rows = []
    with open(gff_path, 'r') as f:
        for line in f:
            if line.startswith('#'):
                continue
            parts = line.rstrip('\n').split('\t')
            if len(parts) < 9:
                continue

            seqid, source, feature, start, end, score, strand, phase, attrs = parts
            if feature != 'CDS':
                continue

            att = parse_gff_attributes(attrs)
            gene_id = att.get(gene_attr) or att.get('locus_tag') or att.get('ID')
            if not gene_id:
                continue

            rows.append({
                'gene': gene_id,
                'contig': seqid,
                'start': int(start),
                'end': int(end),
                'strand': strand,
                'product': att.get('product', 'hypothetical protein'),
                'gene_name': att.get('gene', gene_id)
            })

    return pd.DataFrame(rows)

def parse_kofam_output(path):
    """
    Parser robusto para KOfamScan.
    - Detecta gene/query id aunque la línea empiece con '*' u otro marcador
    - Busca el primer KO tipo Kxxxxx en cualquier columna
    - Deduplica por gen, conservando la mejor puntuación si existe
    """

    import re
    import pandas as pd

    if not os.path.exists(path):
        return pd.DataFrame(columns=['gene', 'ko', 'ko_desc'])

    rows = []

    with open(path, 'r') as f:
        for line in f:
            line = line.strip()

            if not line:
                continue
            if line.startswith('#'):
                continue

            # probar primero TSV; si no, usar split por espacios
            cols = line.split('\t')
            if len(cols) == 1:
                cols = re.split(r'\s+', line)

            if len(cols) < 2:
                continue

            # detectar gene id
            gene_candidate = cols [0].strip()

            # si la primera columna es un marcador (*, -, +), usar la segunda
            if gene_candidate in {'*', '-', '+', '?', '!'} and len(cols) > 1:
                gene_candidate = cols [1].strip()

            gene = normalize_gene_id(gene_candidate)
            if not gene:
                continue

            # buscar KO en cualquier columna
            ko = None
            ko_idx = None
            for i, c in enumerate(cols):
                m = re.search(r'(K\d{5})', str(c))
                if m:
                    ko = m.group(1)
                    ko_idx = i
                    break

            if not ko:
                continue

            # score opcional: tomar el mayor float razonable de la fila
            score = None
            numeric_vals = []
            for c in cols:
                try:
                    numeric_vals.append(float(c))
                except Exception:
                    pass
            if numeric_vals:
                score = max(numeric_vals)

            ## descripción: todo lo que esté después del KO
            if ko_idx is not None and ko_idx + 1 < len(cols):
                ko_desc = ' '.join(
                    [str(x) for x in cols[ko_idx + 1:] if str(x).strip()]
                )

                # limpiar números iniciales típicos de KOfam
                ko_desc = re.sub(
                    r'^\s*[\d\.eE+-]+\s+[\d\.eE+-]+\s+[\d\.eE+-]+\s*',
                    '',
                    ko_desc
                )
            else:
                ko_desc = ''

            rows.append({
                'gene': gene,
                'ko': ko,
                'ko_desc': ko_desc,
                'score': score
            })

    if not rows:
        return pd.DataFrame(columns=['gene', 'ko', 'ko_desc'])

    dfk = pd.DataFrame(rows)

    # ordenar para conservar el mejor score por gen
    if 'score' in dfk.columns:
        dfk = dfk.sort_values(by='score', ascending=False, na_position='last')

    dfk = dfk.drop_duplicates(subset=['gene'], keep='first').copy()

    return dfk[['gene', 'ko', 'ko_desc']]

def normalize_gene_id(x):
    """
    Normaliza IDs de genes para que GFF, FAA, featureCounts y KOfam
    usen el mismo identificador base.
    """
    if x is None:
        return None
    try:
        if pd.isna(x):
            return None
    except Exception:
        pass

    x = str(x).strip()

    # quitar marcadores tipo '*'
    x = x.lstrip('*').strip()

    # tomar solo el primer token si viene con descripción
    x = x.split()[0]

    # patrón clásico tipo OCJKAIMN_00001
    m = re.search(r'([A-Za-z0-9]+_\d{5,})', x)
    if m:
        return m.group(1)

    return x

#HELPER
def get_faa_gene_ids(faa_path):
    """
    Extrae IDs normalizados del FAA actual.
    """
    ids = []
    if not os.path.exists(faa_path):
        return set()

    with open(faa_path, 'r') as f:
        for line in f:
            if line.startswith('>'):
                hdr = line[1:].strip()
                gid = normalize_gene_id(hdr)
                if gid:
                    ids.append(gid)

    return set(ids)


def is_missing_or_empty(x):
    if x is None:
        return True
    try:
        if pd.isna(x):
            return True
    except Exception:
        pass
    if isinstance(x, str) and not x.strip():
        return True
    if isinstance(x, list) and len(x) == 0:
        return True
    return False


def load_featurecounts(path):
    dfc = pd.read_csv(path, sep='\t', comment='#')
    if dfc.empty:
        raise RuntimeError("featureCounts devolvió tabla vacía.")

    count_col = dfc.columns[-1]
    dfc = dfc.rename(columns={
        'Geneid': 'gene',
        'Length': 'length',
        count_col: 'count'
    })

    total_counted = dfc['count'].sum()
    if total_counted == 0:
        raise RuntimeError("no lectures; therefore no RPKM.")

    dfc['RPKM'] = (1e9 * dfc['count']) / (dfc['length'] * total_counted)
    return dfc[['gene', 'length', 'count', 'RPKM']].copy(), total_counted

def fetch_kegg_pathway_names():
    url = 'https://rest.kegg.jp/list/pathway'
    r = requests.get(url, timeout=300)
    r.raise_for_status()

    names = {}
    for line in r.text.strip().split('\n'):
        if '\t' not in line:
            continue
        pid, pname = line.split('\t', 1)
        names[pid.replace('path:', '').strip()] = pname.strip()
    return names

def fetch_kegg_pathways_for_kos_fast(kos):
    """
    Descarga el mapping global KO -> pathway una sola vez y filtra localmente.
    """
    target_kos = set([k for k in kos if pd.notna(k)])
    ko_to_path = {ko: [] for ko in target_kos}

    url = 'https://rest.kegg.jp/link/pathway/ko'
    r = requests.get(url, timeout=600)
    r.raise_for_status()

    for line in r.text.strip().split('\n'):
        if '\t' not in line:
            continue
        left, right = line.split('\t')
        ko = left.replace('ko:', '').strip()
        pw = right.replace('path:', '').strip()
        if ko in target_kos and pw:
            ko_to_path.setdefault(ko, []).append(pw)

    for ko in ko_to_path:
        ko_to_path[ko] = sorted(set(ko_to_path[ko]))

    return ko_to_path

def classify_gene(row):

    txt = f"{row.get('product', '')} {row.get('gene_name', '')} {row.get('ko_desc', '')}".lower()

    # ------------------------------------------------------------------
    # 1) RESISTANCE
    # ------------------------------------------------------------------
    if any(k in txt for k in [

        # Bombas de eflujo
        'rnd efflux',
        'efflux pump',
        'multidrug efflux',
        'drug transporter',
        'macrolide efflux',
        'tetracycline efflux',
        'quinolone efflux',
        'chloramphenicol efflux',

        # β-lactamasas
        'beta-lactamase',
        'beta-lactam',
        'beta lactam',
        'metallo-beta-lactamase',
        'metallobetalactamase',
        'class a beta-lactamase',
        'class b beta-lactamase',
        'class c beta-lactamase',
        'class d beta-lactamase',
        'carbapenemase',
        'extended spectrum beta-lactamase',
        'esbl',
        'ampc',

        # Resistencia por familia
        'aminoglycoside',
        'tetracycline',
        'vancomycin',
        'teicoplanin',
        'colistin',
        'polymyxin',
        'fosfomycin',
        'chloramphenicol',
        'quinolone',
        'fluoroquinolone',
        'ciprofloxacin',
        'macrolide',
        'erythromycin',
        'rifampicin',
        'rifampin',
        'sulfonamide',
        'trimethoprim',
        'linezolid',
        'daptomycin',

        # Genes AMR frecuentes
        'bla1',
        'bla2',
        'blandm',
        'blakpc',
        'blavim',
        'blaimp',
        'blatem',
        'blashv',
        'blaoxa',
        'mcr-1',
        'mcr-2',

        # Generales
        'antibiotic resistance',
        'drug resistance',
        'multidrug resistance',
        'resistance determinant'

    ]):
        return 'Resistance'

    # ------------------------------------------------------------------
    # 2) EXTREMOPHILE-LIKE
    # ------------------------------------------------------------------
    if any(k in txt for k in [

        # Temperatura
        'heat shock',
        'heat shock protein',
        'cold shock',
        'cold shock protein',
        'thermophilic',
        'hyperthermophilic',
        'psychrophilic',

        # Estrés osmótico
        'osmotic stress',
        'osmotic shock',
        'osmoprotectant',
        'osmoprotect',
        'compatible solute',
        'ectoine',
        'hydroxyectoine',
        'glycine betaine',

        # pH
        'acid resistance',
        'acid tolerance',
        'alkaliphilic',
        'acidophilic',
        'ph tolerance',

        # Sal
        'halophilic',
        'halotolerant',
        'extreme halophile',

        # Metales pesados
        'heavy metal tolerance',
        'heavy metal resistance',
        'cadmium resistance',
        'cadmium tolerance',
        'mercury resistance',
        'mercury tolerance',
        'arsenic resistance',
        'arsenic tolerance',
        'copper resistance',
        'zinc resistance',
        'metal detoxification',
        'arsenic operon',
        'mer operon',

        # Estrés general
        'desiccation tolerance',
        'starvation response',
        'stress response',
        'general stress',
        'universal stress'

    ]):
        return 'Extremophile-like'

    # ------------------------------------------------------------------
    # 3) BIOPLASTIC-LIKE
    # ------------------------------------------------------------------
    if any(k in txt for k in [

        'polyhydroxyalkanoate',
        'polyhydroxybutyrate',
        'pha synthase',
        'phb synthase',
        'biopolyester',
        'phaa', 'phab', 'phac',
        'phaz', 'phar',
        'phba', 'phbb', 'phbc',
        'phbz',
        'phaa',
        'phab',
        'phac',
        'pha operon',
        'phb operon'

    ]):
        return 'Bioplastic-like'

    # ------------------------------------------------------------------
    # 4) BIODEGRADATION-LIKE
    # ------------------------------------------------------------------
    if any(k in txt for k in [

        'alkane monooxygenase',
        'alkane hydroxylase',
        'alkb',
        'dioxygenase',
        'catechol dioxygenase',
        'protocatechuate dioxygenase',
        'benzoate degradation',
        'toluene degradation',
        'xylene degradation',
        'phenol degradation',
        'phenol hydroxylase',
        'styrene monooxygenase',
        'biphenyl dioxygenase',
        'naphthalene dioxygenase',
        'petase',
        'mhetase',
        'polyester hydrolase',
        'cutinase',
        'plastic degradation',
        'polyethylene',
        'polypropylene',
        'polyurethane',
        'nylonase',
        'xenobiotic degradation',
        'hydrocarbon degradation',
        'oil degradation',
        'diesel degradation'

    ]):
        return 'Biodegradation-like'

    # ------------------------------------------------------------------
    # 5) INDUSTRIAL-ENZYME-LIKE
    # ------------------------------------------------------------------
    if any(k in txt for k in [

        # CAZymes
        'alpha-amylase',
        'beta-amylase',
        'amylase',
        'pullulanase',
        'xylanase',
        'cellulase',
        'pectinase',
        'pectate lyase',
        'chitinase',
        'beta-glucosidase',
        'beta-galactosidase',

        # Proteasas
        'subtilisin',
        'protease',
        'serine protease',
        'keratinase',
        'collagenase',

        # Lipasas
        'lipase',
        'esterase',
        'carboxylesterase',

        # Oxidasas industriales
        'laccase',
        'tyrosinase',
        'peroxidase'
    ]):
        return 'Industrial-enzyme-like'

    # ------------------------------------------------------------------
    # 6) VIRULENCE
    # ------------------------------------------------------------------
    if any(k in txt for k in [

        'virulence',
        'virulence factor',
        'toxin',
        'exotoxin',
        'endotoxin',
        'enterotoxin',
        'hemolysin',
        'leukocidin',
        'adhesin',
        'adhesion',
        'fimbr',
        'fimbriae',
        'pilus',
        'curli',
        'capsule',
        'capsular polysaccharide',
        'invasion',
        'invasin',
        'internalin',
        'secretion system',
        'type iii secretion',
        'type iv secretion',
        'type vi secretion',
        't3ss',
        't4ss',
        't6ss',
        'siderophore',
        'iron acquisition',
        'pathogenic',
        'pathogenicity'

    ]):
        return 'Virulence'

    return 'Other'

def build_kegg_gene_graph(df_input):
    G = nx.Graph()

    for _, row in df_input.iterrows():
        gene = row['gene']
        G.add_node(
            gene,
            rpkm=float(row['RPKM']),
            count=int(row['count']),
            category=row['cat'],
            product=row.get('product', 'NA'),
            ko=row.get('ko', None),
            pathways=row.get('pathways', []),
            core_like=bool(row.get('core_like', False))
        )

    genes = df_input['gene'].tolist()

    for i in range(len(genes)):
        g1 = genes[i]
        p1 = set(G.nodes[g1]['pathways']) if G.nodes[g1]['pathways'] else set()

        for j in range(i + 1, len(genes)):
            g2 = genes[j]
            p2 = set(G.nodes[g2]['pathways']) if G.nodes[g2]['pathways'] else set()
            shared = p1.intersection(p2)
            if shared:
                G.add_edge(g1, g2, weight=len(shared), shared_pathways=list(shared))

    return G

# ==============================================================================
# UTILIDADES PARA REFERENCIAS EXTERNAS CORE-LIKE
# ==============================================================================

import re
import pandas as pd

CORE_REF_SOURCES = {
    "EscherichiaShigella": os.path.join(CORE_REF_DIR, "cgEscherichiaShigella.txt"),
    "Salmonella": os.path.join(CORE_REF_DIR, "cgSalmonella.txt"),
    "Yersinia": os.path.join(CORE_REF_DIR, "cgYersinia.txt"),
    "Clostridioides": os.path.join(CORE_REF_DIR, "cgClostridiioides.txt"),
    "MycobacteriumTB": os.path.join(CORE_REF_DIR, "cgMycobacteriumTB.txt")
}

def normalize_text(x):
    if pd.isna(x) or x is None:
        return ""
    x = str(x).strip().lower()
    x = re.sub(r"\s+", " ", x)
    return x

def parse_locus_tag_field(x):
    """
    Ejemplos:
      'b2231 (gyrA)' -> locus_id='b2231', gene_symbol='gyrA'
      'CD630_00010(dnaA)' -> locus_id='CD630_00010', gene_symbol='dnaA'
      'MTBC-00001 (Rv2925c,rnc)' -> locus_id='MTBC-00001', gene_symbol='rnc'
      'YPO0679' -> locus_id='YPO0679', gene_symbol=''
    """
    if pd.isna(x):
        return "", ""

    raw = str(x).strip()

    m = re.match(r"^(.*?)\((.*?)\)\s*$", raw)
    if m:
        locus_id = m.group(1).strip()
        inside = m.group(2).strip()

        parts = [p.strip() for p in inside.split(",") if p.strip()]
        gene_symbol = parts[-1] if parts else ""
        return locus_id, gene_symbol

    return raw.strip(), ""

def load_core_reference(ref_name):
    source = CORE_REF_SOURCES.get(ref_name)
    if source is None:
        raise ValueError(f"Referencia no reconocida: {ref_name}")
    if not os.path.exists(source):
        raise FileNotFoundError(f"No se encontró el archivo de referencia: {source}")

    df_ref = pd.read_csv(source, sep="\t")
    df_ref.columns = [c.strip() for c in df_ref.columns]

    if 'Locus Tag' not in df_ref.columns:
        raise ValueError(f"El archivo {source} no contiene la columna 'Locus Tag'")
    if 'Description' not in df_ref.columns:
        df_ref['Description'] = ""

    parsed = df_ref['Locus Tag'].apply(parse_locus_tag_field)
    df_ref['locus_id'] = parsed.apply(lambda x: x[0])
    df_ref['gene_symbol'] = parsed.apply(lambda x: x[1])

    df_ref['locus_id_norm'] = df_ref['locus_id'].apply(normalize_text)
    df_ref['gene_symbol_norm'] = df_ref['gene_symbol'].apply(normalize_text)
    df_ref['description_norm'] = df_ref['Description'].apply(normalize_text)

    return df_ref

def match_against_core_reference(row, df_ref):
    """
    Devuelve:
      matched (bool)
      match_type (str)
      matched_ref_locus (str)
      matched_ref_gene (str)
      matched_ref_description (str)
    """
    row_gene = normalize_text(row.get('gene', ''))
    row_gene_name = normalize_text(row.get('gene_name', ''))
    row_product = normalize_text(row.get('product', ''))
    row_ko_desc = normalize_text(row.get('ko_desc', ''))

    searchable = " | ".join([row_gene, row_gene_name, row_product, row_ko_desc])

    # 1) exact match por locus_id
    if row_gene:
        hit = df_ref[df_ref['locus_id_norm'] == row_gene]
        if not hit.empty:
            h = hit.iloc[0]
            return True, "exact_locus_id", h['locus_id'], h['gene_symbol'], h['Description']

    # 2) exact match por gene_symbol
    if row_gene_name:
        hit = df_ref[df_ref['gene_symbol_norm'] == row_gene_name]
        if not hit.empty:
            h = hit.iloc[0]
            return True, "exact_gene_symbol", h['locus_id'], h['gene_symbol'], h['Description']

    # 3) partial match por gene_symbol dentro de gene/product/ko_desc
    for _, h in df_ref.iterrows():
        gs = h['gene_symbol_norm']
        if gs and gs in searchable:
            return True, "partial_gene_symbol", h['locus_id'], h['gene_symbol'], h['Description']

    # 4) partial match por descripción
    for _, h in df_ref.iterrows():
        desc = h['description_norm']
        if desc and len(desc) >= 10 and desc in searchable:
            return True, "partial_description", h['locus_id'], h['gene_symbol'], h['Description']

    return False, "", "", "", ""

# ==============================================================================
# 6) INPUT DEL USUARIO
# ==============================================================================

print("\nWELCOME TO SELECT-NET for Bacteria")

# ----------------------------------------------------------------------
# Origen de datos
# ----------------------------------------------------------------------
print("\nData Origin:")
print(" [S] SRA/ENA accession")
print(" [L] FASTQ local from Google Drive")

input_mode = ""
while input_mode not in ['S', 'L']:
    input_mode = input("Choose S o L: ").strip().upper()

sample_id = None
local_r1 = None
local_r2 = None

if input_mode == 'S':
    sample_id = input("SRA ID (ex. SRR920204): ").strip()
else:
    local_r1 = input("Path in Drive to R1 FASTQ/FASTQ.GZ: ").strip()
    local_r2 = input("Path in Drive to R2 FASTQ/FASTQ.GZ (or just press ENTER if its single-end): ").strip()
    if local_r2 == "":
        local_r2 = None

    # construir identificador amigable a partir del nombre del archivo
    base_name = Path(local_r1).name

    # quitar extensiones frecuentes
    for ext in ['.fastq.gz', '.fq.gz', '.fastq', '.fq']:
        if base_name.endswith(ext):
            base_name = base_name[:-len(ext)]
            break

    # quitar sufijos comunes de R1
    for suffix in ['_R1', '_1', '.R1', '.1']:
        if base_name.endswith(suffix):
            base_name = base_name[:-len(suffix)]
            break

    sample_id = base_name

# ----------------------------------------------------------------------
# Tipo de biblioteca
# ----------------------------------------------------------------------
print("\n Library Type:")
print(" [R] RNA-seq / metatranscriptome")
print(" [G] WGS / metagenome")

lib_mode = ""
while lib_mode not in ['R', 'G']:
    lib_mode = input("Select R o G: ").strip().upper()

print("\nAssembly type:")
print(" [G] Genome (isolate)")
print(" [M] Metagenome")

assembly_type = ""
while assembly_type not in ['G', 'M']:
    assembly_type = input("Select G o M: ").strip().upper()

# ---------------------------------------------------------------------
# Analyses Mode
# ---------------------------------------------------------------------
print("\n--- MODES ---")
print(" [A] Thematic (extremophile + biotech-like)")
print(" [B] Quantitative")
print(" [C] Comparative-like (remove all core-like)")
print(" [CH] Comparative-high (remove core-like externos + highly expressed)")
print(" [CL] Comparative-low (remove core-like externos + lowly expressed)")
print(" [D] Epidemiological")
print("-------------")

mode = ""
while mode not in ['A', 'B', 'C', 'CH', 'CL', 'D']:
    mode = input("Select mode (A/B/C/CH/CL/D): ").strip().upper()

# -----------------------------------------------------------------------------
# REFERENCIA EXTERNA CORE-LIKE (solo si el modo es C/CH/CL)
# -----------------------------------------------------------------------------
external_core_ref = None
external_core_ref_applied_to_filter = False

if mode in ['C', 'CH', 'CL']:
    print("\nExternal core-like (to C/CH/CL):")
    print(" [1] EscherichiaShigella")
    print(" [2] Salmonella")
    print(" [3] Yersinia")
    print(" [4] Clostridioides")
    print(" [5] MycobacteriumTB")

    ext_choice = ""
    while ext_choice not in ['1', '2', '3', '4', '5']:
        ext_choice = input("Select external-core (1-5): ").strip()

    EXT_REF_MAP = {
        '1': "EscherichiaShigella",
        '2': "Salmonella",
        '3': "Yersinia",
        '4': "Clostridioides",
        '5': "MycobacteriumTB"
    }

    external_core_ref = EXT_REF_MAP[ext_choice]
    external_core_ref_applied_to_filter = True

print(f"external_core_ref: {external_core_ref}")
print(f"external_core_ref_applied_to_filter: {external_core_ref_applied_to_filter}")

# ==============================================================================
# 2) WORKSPACE
# ==============================================================================

try:
    from google.colab import drive
    drive.mount('/content/drive')
    DRIVE_BASE = '/content/drive/MyDrive/SELECT-NET'
except Exception:
    print("Google Drive no disponible. Se usará salida local en /content/SELECT-NET_output")
    DRIVE_BASE = '/content/SELECT-NET_output'

OUT_DIR = os.path.join(DRIVE_BASE, "output")

from datetime import datetime

CHECKPOINT_DIR = os.path.join(DRIVE_BASE, sample_id, "checkpoints")

os.makedirs(CHECKPOINT_DIR, exist_ok=True)

WORK_DIR = '/content/selectnet_work'
CACHE_DIR = '/content/selectnet_cache'
TMP_DIR = '/content/selectnet_tmp'

ensure_dir(OUT_DIR)
ensure_dir(WORK_DIR)
ensure_dir(CACHE_DIR)
ensure_dir(TMP_DIR)

os.chdir(WORK_DIR)

print(f"WORK_DIR (local): {WORK_DIR}")
print(f"CACHE_DIR (local): {CACHE_DIR}")
print(f"OUT_DIR (Drive): {OUT_DIR}")
print(f"CHECKPOINT_DIR (Drive): {CHECKPOINT_DIR}")
# ==============================================================================
# Looking CHECKPOINTS
# ==============================================================================

if os.path.exists(CHECKPOINT_DIR) and len(os.listdir(CHECKPOINT_DIR)) > 0:
    print("Checkpoints detected → reusing existing information")
else:
    print("No checkpoints founds → Full pipeline will be implemented")

# ==============================================================================
# 2A) CHECKPOINT HELPERS
# ==============================================================================

def cp_to_checkpoint(src, name=None):
    if not src or not os.path.exists(src):
        return None
    dst = os.path.join(CHECKPOINT_DIR, name if name else os.path.basename(src))
    shutil.copy2(src, dst)
    print(f"Checkpoint save: {dst}")
    return dst

def cpdir_to_checkpoint(src_dir, name=None):
    if not src_dir or not os.path.exists(src_dir):
        return None
    dst = os.path.join(CHECKPOINT_DIR, name if name else os.path.basename(src_dir))
    if os.path.exists(dst):
        shutil.rmtree(dst)
    shutil.copytree(src_dir, dst)
    print(f"Checkpoint save2: {dst}")
    return dst

def restore_from_checkpoint(name, dst_path):
    src = os.path.join(CHECKPOINT_DIR, name)
    if not os.path.exists(src):
        return False
    parent = os.path.dirname(dst_path)
    if parent:
        ensure_dir(parent)
    if os.path.isdir(src):
        if os.path.exists(dst_path):
            shutil.rmtree(dst_path)
        shutil.copytree(src, dst_path)
    else:
        shutil.copy2(src, dst_path)
    print(f"Restored from checkpoint: {src} -> {dst_path}")
    return True

def checkpoint_exists(name):
    return os.path.exists(os.path.join(CHECKPOINT_DIR, name))

# ==============================================================================
# 3) INSTALACIÓN DE HERRAMIENTAS
# ==============================================================================

print("\nInstalling system tools...")
run(['apt-get', 'update', '-qq'])

run([
    'apt-get', 'install', '-y', '-qq',
    'samtools', 'bedtools', 'megahit', 'minimap2',
    'aria2', 'cutadapt', 'fastqc', 'trim-galore', 'wget',
    'openjdk-17-jre-headless',
    'bioperl',
    'libbio-searchio-hmmer-perl',
    'ncbi-blast+',
    'ncbi-tools-bin',
    'prodigal', 'hmmer', 'parallel', 'infernal',
    'aragorn',
    'ruby', 'git', 'subread', 'less'
])

run([
    'apt-get', 'install', '-y', '-qq',
    'spades'
])

print("Installing Python...")
run([
    sys.executable, '-m', 'pip', 'install', '-q', '--upgrade',
    'pandas', 'networkx', 'matplotlib', 'requests', 'biopython'
])

# ==============================================================================
# 3A) ASEGURAR CUTADAPT
# ==============================================================================
print("\nInstalling Cutadapt on Python...")
run([sys.executable, '-m', 'pip', 'install', '--upgrade', 'cutadapt', '-q'])

CUTADAPT_BIN = shutil.which('cutadapt')
if not CUTADAPT_BIN:
    raise RuntimeError("not ex at cutadapt en PATH after install.")

print(f"Cutadapt detected: {CUTADAPT_BIN}")

for cmd in [[CUTADAPT_BIN, '--version'], ['trim_galore', '--version'], ['fastqc', '--version']]:
    res = subprocess.run(cmd, capture_output=True, text=True, check=False)
    print(f"\n[VERSION] {' '.join(cmd)}")
    print(res.stdout if res.stdout else res.stderr)


# ==============================================================================
# 3B) PROKKA DESDE REPOSITORIO (NO apt)
# ==============================================================================
print("\nPreparing Prokka from repository...")

aragorn_bin = shutil.which('aragorn')
if not aragorn_bin:
    raise RuntimeError("'aragorn' not found at PATH. Check system instalation.")
print(f"Aragorn detectado en: {aragorn_bin}")

tbl2asn_bin = shutil.which('tbl2asn')
if not tbl2asn_bin:
    raise RuntimeError("'tbl2asn' not found at PATH. Check installation at ncbi-tools-bin.")
print(f"tbl2asn detectado en: {tbl2asn_bin}")

PROKKA_DIR = '/content/prokka'

if os.path.exists(PROKKA_DIR):
    shutil.rmtree(PROKKA_DIR)

run(['git', 'clone', '--depth', '1', 'https://github.com/tseemann/prokka.git', PROKKA_DIR])

PROKKA_BIN = os.path.join(PROKKA_DIR, 'bin', 'prokka')
if not os.path.exists(PROKKA_BIN):
    raise RuntimeError("No se encontró /content/prokka/bin/prokka tras clonar el repositorio.")

print(f"Prokka detected en: {PROKKA_BIN}")

aragorn_bin = shutil.which('aragorn')
if not aragorn_bin:
    raise RuntimeError("'aragorn' not found at PATH.")
print(f"Aragorn detected: {aragorn_bin}")

tbl2asn_bin = shutil.which('tbl2asn')
if not tbl2asn_bin:
    raise RuntimeError("'tbl2asn' not found at PATH.")
print(f"tbl2asn detected at: {tbl2asn_bin}")

check_hmmer_perl = subprocess.run(
    ['perl', '-MBio::SearchIO::hmmer3', '-e', 'print "OK\n"'],
    capture_output=True, text=True, check=False
)

if check_hmmer_perl.returncode != 0:
    raise RuntimeError(
        "no loaded Bio::SearchIO::hmmer3. "
        "check installation of libbio-searchio-hmmer-perl."
    )

print("Bio::SearchIO::hmmer3 disponible")

# ==============================================================================
# 3C) TBL2ASN
# ==============================================================================

print("\nVerifying tbl2asn...")

tbl2asn_bin = shutil.which('tbl2asn')
if not tbl2asn_bin:
    raise RuntimeError(
        "'tbl2asn' not at PATH. "
        "Install the package 'ncbi-tools-bin' beffore ex Prokka."
    )

print(f"tbl2asn detected at: {tbl2asn_bin}")


# ==============================================================================
# 3D) SETUP DB DE PROKKA
# ==============================================================================
setup_res = subprocess.run(
    [PROKKA_BIN, '--setupdb'],
    capture_output=True,
    text=True,
    check=False
)

print("\n---- STDOUT prokka --setupdb ----")
print(setup_res.stdout if setup_res.stdout else "[vacío]")

print("\n---- STDERR prokka --setupdb ----")
print(setup_res.stderr if setup_res.stderr else "[vacío]")

if setup_res.returncode != 0:
    raise RuntimeError(
        f"prokka --setupdb fail with code {setup_res.returncode}. "
        "Not possible a trustable annotation."
    )

listdb_res = subprocess.run(
    [PROKKA_BIN, '--listdb'],
    capture_output=True,
    text=True,
    check=False
)

print("\n[PROKKA CHECK]", PROKKA_BIN, "--listdb")
print(listdb_res.stdout if listdb_res.stdout else listdb_res.stderr)

# ==============================================================================
# 4) SRA TOOLKIT
# ==============================================================================

print("\nPreparing SRA Toolkit...")
os.chdir(CACHE_DIR)

sra_tar = os.path.join(CACHE_DIR, 'sra.tar.gz')
if not os.path.exists(sra_tar):
    run([
        'wget', '-q', '-O', sra_tar,
        'https://ftp-trace.ncbi.nlm.nih.gov/sra/sdk/current/sratoolkit.current-ubuntu64.tar.gz'
    ])

if not glob.glob('sratoolkit.*/bin'):
    run(['tar', '-xzf', 'sra.tar.gz'])

sra_bins = glob.glob(os.path.join(CACHE_DIR, 'sratoolkit.*/bin'))
if not sra_bins:
    raise RuntimeError("No se detectó SRA Toolkit.")

os.environ['PATH'] = f"{os.path.abspath(sra_bins[0])}:{os.environ.get('PATH', '')}"
run(['vdb-config', '--restore-defaults'], check=False)
run(['vdb-config', '--prefetch-to-cwd'], check=False)

os.chdir(WORK_DIR)
print("SRA Toolkit ready at PATH.")

# ==============================================================================
# 5) KOFAMSCAN
# ==============================================================================

print("\nPreparing KOfamScan...")
KOFAM_SCAN_DIR = '/content/kofam_scan'
KOFAM_DB_DIR = '/content/kofam_db'

if os.path.exists(KOFAM_SCAN_DIR):
    shutil.rmtree(KOFAM_SCAN_DIR)

run(['git', 'clone', '--depth', '1', 'https://github.com/takaram/kofam_scan.git', KOFAM_SCAN_DIR])

ensure_dir(KOFAM_DB_DIR)
os.chdir(KOFAM_DB_DIR)

profiles_tar = os.path.join(KOFAM_DB_DIR, 'profiles.tar.gz')
ko_list_gz = os.path.join(KOFAM_DB_DIR, 'ko_list.gz')

if not os.path.exists(profiles_tar):
    run(['wget', '-q', '-O', profiles_tar, 'ftp://ftp.genome.jp/pub/db/kofam/profiles.tar.gz'], check=False)

if not os.path.exists(ko_list_gz):
    run(['wget', '-q', '-O', ko_list_gz, 'ftp://ftp.genome.jp/pub/db/kofam/ko_list.gz'], check=False)

if os.path.exists(profiles_tar) and os.path.getsize(profiles_tar) > 0:
    run(['tar', '-xzf', profiles_tar], check=False)

if os.path.exists(ko_list_gz) and os.path.getsize(ko_list_gz) > 0:
    run(['gunzip', '-f', ko_list_gz], check=False)

os.chdir(WORK_DIR)

# ==============================================================================
# 7) DESCARGA DE READS
# ==============================================================================

print(f"\nPreparing reads {sample_id}...")
r1, r2 = None, None

def download_from_ena(sample_id):
    ena_url = (
        f"https://www.ebi.ac.uk/ena/portal/api/filereport"
        f"?accession={sample_id}&result=read_run&fields=fastq_ftp&format=tsv"
    )
    resp = requests.get(ena_url, timeout=60)
    resp.raise_for_status()

    lines = resp.text.strip().split('\n')
    if len(lines) < 2:
        return []

    fields = lines[1].split('\t')
    if not fields or not fields[0].strip():
        return []

    ftp_files = fields[0].split(';')
    downloaded = []

    for ftp in ftp_files:
        ftp = ftp.strip()
        if not ftp:
            continue

        if ftp.startswith("ftp."):
            url = "https://" + ftp
        elif ftp.startswith("http://") or ftp.startswith("https://"):
            url = ftp
        else:
            url = "https://" + ftp

        out_name = os.path.join(WORK_DIR, os.path.basename(ftp))
        run(['wget', '-q', '-O', out_name, url])
        downloaded.append(out_name)

    return downloaded


if input_mode == 'S':
    print(f"Downloading reads {sample_id} from SRA/ENA...")

    try:
        run(['prefetch', sample_id])
        run([
            'fasterq-dump',
            '--split-3',
            '--force',
            '--threads', '4',
            '--temp', TMP_DIR,
            '-O', WORK_DIR,
            sample_id
        ])

        candidate_r1 = first_existing([
            os.path.join(WORK_DIR, f'{sample_id}_1.fastq'),
            os.path.join(WORK_DIR, f'{sample_id}.fastq')
        ])

        candidate_r2 = first_existing([
            os.path.join(WORK_DIR, f'{sample_id}_2.fastq')
        ])

        r1 = candidate_r1
        r2 = candidate_r2

    except Exception as e:
        print(f"SRA fail ({e}). Trying ENA...")
        ena_files = download_from_ena(sample_id)
        if len(ena_files) >= 1:
            r1 = ena_files[0]
        if len(ena_files) >= 2:
            r2 = ena_files[1]

elif input_mode == 'L':
    print("Using FASTQ local from Google Drive...")

    if not os.path.exists(local_r1):
        raise RuntimeError(f"R1 No exist at Drive: {local_r1}")

    r1 = os.path.join(WORK_DIR, os.path.basename(local_r1))
    shutil.copy2(local_r1, r1)

    if local_r2:
        if not os.path.exists(local_r2):
            raise RuntimeError(f"R2 No exist at Drive: {local_r2}")
        r2 = os.path.join(WORK_DIR, os.path.basename(local_r2))
        shutil.copy2(local_r2, r2)

if not r1:
    raise RuntimeError("Any FASTQ could be recovered.")

print(f"Inputs detected: r1={r1}, r2={r2}")

# ==============================================================================
# 8) TRIMMING CON TRIM GALORE + CUTADAPT EXPLÍCITO
# ==============================================================================

if not FORCE_RERUN and SAVE_TRIM_CHECKPOINTS and checkpoint_exists('clean_r1.fastq'):
    restore_from_checkpoint('clean_r1.fastq', 'clean_r1.fastq')
    clean_r1 = 'clean_r1.fastq'

    if checkpoint_exists('clean_r2.fastq'):
        restore_from_checkpoint('clean_r2.fastq', 'clean_r2.fastq')
        clean_r2 = 'clean_r2.fastq'
    else:
        clean_r2 = None

    print(f"Restored Trimming from checkpoint: {clean_r1}, {clean_r2}")

else:
    print("\nEx Trim-Galore...")

    if not CUTADAPT_BIN:
        raise RuntimeError("CUTADAPT_BIN not defined. Not Trim Galore running.")

    for f in [r1, r2]:
        if f:
            print(f"\n[CHECK FILE] {f}")
            print("Exist:", os.path.exists(f))
            if os.path.exists(f):
                print("Size (bytes):", os.path.getsize(f))

    if r2:
        trim_cmd = [
            'trim_galore',
            '--paired',
            '--path_to_cutadapt', CUTADAPT_BIN,
            '-q', '20',
            '--length', '36',
            '--output_dir', '.',
            r1, r2
        ]
    else:
        trim_cmd = [
            'trim_galore',
            '--path_to_cutadapt', CUTADAPT_BIN,
            '-q', '20',
            '--length', '36',
            '--output_dir', '.',
            r1
        ]

    print("\n[TRIM CMD]")
    print(' '.join(trim_cmd))

    trim_res = subprocess.run(trim_cmd, capture_output=True, text=True, check=False)

    print("\n---- STDOUT Trim Galore ----")
    print(trim_res.stdout if trim_res.stdout else "[vacío]")

    print("\n---- STDERR Trim Galore ----")
    print(trim_res.stderr if trim_res.stderr else "[vacío]")

    if trim_res.returncode != 0:
        raise RuntimeError(f"Trim Galore falló con código {trim_res.returncode}. Revisa STDOUT/STDERR impresos arriba.")

    base1 = os.path.splitext(os.path.basename(r1))[0]

    if r2:
        base2 = os.path.splitext(os.path.basename(r2))[0]

        src1 = first_existing([
            f"{base1}_val_1.fq",
            f"{base1}_val_1.fq.gz"
        ] + glob.glob("*_val_1.f*"))

        src2 = first_existing([
            f"{base2}_val_2.fq",
            f"{base2}_val_2.fq.gz"
        ] + glob.glob("*_val_2.f*"))

        if not src1 or not src2:
            raise FileNotFoundError("No paired-end clean reads found on Trim Galore.")

        if src1.endswith('.gz'):
            run(['gunzip', '-f', src1])
            src1 = src1[:-3]

        if src2.endswith('.gz'):
            run(['gunzip', '-f', src2])
            src2 = src2[:-3]

        if os.path.exists('clean_r1.fastq'):
            os.remove('clean_r1.fastq')
        if os.path.exists('clean_r2.fastq'):
            os.remove('clean_r2.fastq')

        os.rename(src1, 'clean_r1.fastq')
        os.rename(src2, 'clean_r2.fastq')
        clean_r1, clean_r2 = 'clean_r1.fastq', 'clean_r2.fastq'

    else:
        src = first_existing([
            f"{base1}_trimmed.fq",
            f"{base1}_trimmed.fq.gz"
        ] + glob.glob("*_trimmed.f*"))

        if not src:
            raise FileNotFoundError("not single-end clean reads on Trim Galore.")

        if src.endswith('.gz'):
            run(['gunzip', '-f', src])
            src = src[:-3]

        if os.path.exists('clean_r1.fastq'):
            os.remove('clean_r1.fastq')

        os.rename(src, 'clean_r1.fastq')
        clean_r1, clean_r2 = 'clean_r1.fastq', None

    print(f"Reads limpias: {clean_r1}, {clean_r2}")

    # FastQC separado
    fastqc_inputs = [clean_r1] + ([clean_r2] if clean_r2 else [])
    fastqc_cmd = ['fastqc', '-o', '.', *fastqc_inputs]

    print("\n[FASTQC CMD]")
    print(' '.join(fastqc_cmd))

    fastqc_res = subprocess.run(fastqc_cmd, capture_output=True, text=True, check=False)

    print("\n---- STDOUT FastQC ----")
    print(fastqc_res.stdout if fastqc_res.stdout else "[vacío]")

    print("\n---- STDERR FastQC ----")
    print(fastqc_res.stderr if fastqc_res.stderr else "[vacío]")

    if fastqc_res.returncode != 0:
        print("FastQC fail, but trimming happens and pipeline will continue.")

    if SAVE_TRIM_CHECKPOINTS:
        cp_to_checkpoint(clean_r1, 'clean_r1.fastq')
        if clean_r2:
            cp_to_checkpoint(clean_r2, 'clean_r2.fastq')

# ==============================================================================
# 9) ENSAMBLE
# ==============================================================================

contigs_file = 'assembly_out/final.contigs.fa'

use_existing_assembly = False
use_existing_annotation = False
# ==============================================================================
# CHECKPOINT DE ENSAMBLE
# ==============================================================================

if not FORCE_RERUN and checkpoint_exists('final.contigs.sanitized.fa'):

    ensure_dir('assembly_out')

    restore_from_checkpoint(
        'final.contigs.sanitized.fa',
        'assembly_out/final.contigs.sanitized.fa'
    )

    contigs_file = 'assembly_out/final.contigs.sanitized.fa'

    print("Assembly restored from checkpoint.")

    info_file = os.path.join(
        CHECKPOINT_DIR,
        "checkpoint_info.json"
    )

    if os.path.exists(info_file):

        with open(info_file, "r") as f:
            checkpoint_info = json.load(f)

        print("\nCHECKPOINT INFO")
        print("=" * 50)

        for k, v in checkpoint_info.items():
            print(f"{k}: {v}")

        print("=" * 50)

    else:

        print("checkpoint_info.json not found")

    use_existing_annotation = False


else:

    # ==============================================================================
    # DETECTAR ENSAMBLAJE + ANOTACIÓN EXISTENTE
    # ==============================================================================

    use_existing_annotation = False
    use_existing_assembly = False

    existing_contigs = None
    existing_gff = None
    existing_faa = None

    input_dir = os.path.join(DRIVE_BASE, "input")

    if os.path.exists(input_dir):

        contig_candidates = list(Path(input_dir).glob("*.fna"))
        gff_candidates = list(Path(input_dir).glob("*.gff"))
        faa_candidates = list(Path(input_dir).glob("*.faa"))

        print(f"Searching contigs at: {input_dir}")
        print(f"FNA: {contig_candidates}")
        print(f"GFF: {gff_candidates}")
        print(f"FAA: {faa_candidates}")

        if (
            len(contig_candidates) > 0 and
            len(gff_candidates) > 0 and
            len(faa_candidates) > 0
        ):

            existing_contigs = str(contig_candidates[0])
            existing_gff = str(gff_candidates[0])
            existing_faa = str(faa_candidates[0])

            use_existing_assembly = True

    # ==============================================================================
    # Action
    # ==============================================================================

    if use_existing_assembly:

        print("Using existing assembly + annotation")

        contigs_file = existing_contigs

        use_existing_annotation = True

        gff_file = existing_gff
        faa_file = existing_faa

        annotation_mode = "external_annotation"
        fc_group_attr = "ID"

    else:

        print("\nEx Assembly...")

        if os.path.exists('assembly_out'):
            shutil.rmtree('assembly_out')

        # ==========================================================
        # SPAdes-fast o metaSPAdes-fast
        # ==========================================================

        if assembly_type == 'G':

            print("Using SPAdes-fast (bacterial genoma)")

            spades_cmd = [
                'spades.py',
                '--only-assembler',
                '-k', '21,33,55',
                '-1', clean_r1,
                '-2', clean_r2,
                '-o', 'assembly_out',
                '--threads', '10',
                '--memory', '10'
            ]

        elif assembly_type == 'M':

            print("Using metaSPAdes-fast (metagenome)")

            spades_cmd = [
                'spades.py',
                '--only-assembler',
                '-k', '21,33,55',
                '--meta',
                '-1', clean_r1,
                '-2', clean_r2,
                '-o', 'assembly_out',
                '--threads', '10',
                '--memory', '10'
            ]

        print("[SPAdes CMD]")
        print(' '.join(spades_cmd))

        start_spades = datetime.now()

        print(
            f"Starting SPAdes: {start_spades}"
        )

        run(spades_cmd)

        print("\nassembly_out:")

        if os.path.exists("assembly_out"):

            for f in sorted(glob.glob("assembly_out/*")):
                print(f)

        else:

            print("assembly_out not found")

        contigs_file = 'assembly_out/contigs.fasta'

        end_spades = datetime.now()

        print(
            f"SPAdes ends: {end_spades}"
        )

        print(
            f"total SPAdes TIME: {end_spades - start_spades}"
        )

        if not os.path.exists(contigs_file):
            raise RuntimeError(
                "No contigs.fasta founds (SPAdes fail)"
            )

        print("Ensamble SPAdes READY.")

        sanitized_contigs = 'assembly_out/contigs.sanitized.fasta'

        sanitize_fasta_headers(
            contigs_file,
            sanitized_contigs
        )

        contigs_file = sanitized_contigs

        cp_to_checkpoint(
            contigs_file,
            'final.contigs.sanitized.fa'
        )

        checkpoint_info = {
            "sample_id": sample_id,
            "assembler": "SPAdes-fast",
            "selectnet_version": "1.0",
            "pathway_filter": True,
            "ko_audit": True,
            "node_labels": True,
            "date": str(datetime.now()),
            "contigs_file": "final.contigs.sanitized.fa"
        }

        with open(
            os.path.join(
                CHECKPOINT_DIR,
                "checkpoint_info.json"
            ),
            "w"
        ) as f:

            json.dump(
                checkpoint_info,
                f,
                indent=4
            )

# ==============================================================================
# 10) ANOTATION (PROKKA OBLIGATORIO + CHECKPOINT)
# ==============================================================================

def annotate_contigs_with_prokka(contigs_file, prokka_bin):

    print("Ex Prokka...")

    outdir = "prokka_out"

    if os.path.exists(outdir):
        shutil.rmtree(outdir)

    cmd = [
        prokka_bin,
        '--outdir', outdir,
        '--prefix', 'sample',
        '--cpus', '4',
        '--kingdom', 'Bacteria',
        '--force',
        contigs_file
    ]

    print("[PROKKA CMD]")
    print(' '.join(cmd))

    run(cmd)

    gff_file = os.path.join(outdir, "sample.gff")
    faa_file = os.path.join(outdir, "sample.faa")

    if not os.path.exists(gff_file):
        raise RuntimeError("Prokka not produce GFF")

    cp_to_checkpoint(gff_file, 'sample.gff')
    cp_to_checkpoint(faa_file, 'sample.faa')

    return gff_file, faa_file, 'prokka', 'ID'
if (
    not FORCE_RERUN and
    checkpoint_exists('sample.gff') and
    checkpoint_exists('sample.faa')
):

    restore_from_checkpoint(
        'sample.gff',
        'sample.gff'
    )

    restore_from_checkpoint(
        'sample.faa',
        'sample.faa'
    )

    gff_file = 'sample.gff'
    faa_file = 'sample.faa'

    annotation_mode = 'prokka'
    fc_group_attr = 'ID'

    use_existing_annotation = True

    print("Prokka restored from checkpoint.")

# EX PROKKA
if not use_existing_annotation:
    gff_file, faa_file, annotation_mode, fc_group_attr = annotate_contigs_with_prokka(
        contigs_file,
        PROKKA_BIN
    )

print(f"Used GFF: {gff_file}")

# ==============================================================================
# 11) MAPPING
# ==============================================================================

print("\nMapping reads to contigs...")

if clean_r2:
    run(['minimap2', '-ax', 'sr', contigs_file, clean_r1, clean_r2, '-o', 'alignments.sam'])
else:
    run(['minimap2', '-ax', 'sr', contigs_file, clean_r1, '-o', 'alignments.sam'])

run('samtools view -bS alignments.sam > alignments.bam', shell=True)
run(['samtools', 'sort', 'alignments.bam', '-o', 'alignments.sorted.bam'])
run(['samtools', 'index', 'alignments.sorted.bam'])

print("BAM ordered and index.")

if SAVE_BAM_CHECKPOINTS:
    cp_to_checkpoint('alignments.sorted.bam', 'alignments.sorted.bam')
    cp_to_checkpoint('alignments.sorted.bam.bai', 'alignments.sorted.bam.bai')

if CLEANUP_AFTER_SUCCESS:
    cleanup_large_files(['alignments.sam', 'alignments.bam'])

# ==============================================================================
# 12) CONTING + RPKM
# ==============================================================================

fc_out = 'gene_counts.txt'

fc_group_attr = "ID"   # definido aquí

if not FORCE_RERUN and checkpoint_exists('gene_counts.txt'):
    restore_from_checkpoint('gene_counts.txt', fc_out)
    df_counts, total_counted = load_featurecounts(fc_out)
    print(f"gene_counts.txt restaurado desde checkpoint. Total contado = {total_counted}")
else:
    print("\nEx featureCounts to RPKM...")


    fc_group_attr = "ID"

    fc_cmd = [
    'featureCounts',
    '-a', gff_file,
    '-F', 'GTF',
    '-t', 'CDS',
    '-g', fc_group_attr,
    '-o', fc_out,
    '-p',
    '--countReadPairs',
    'alignments.sorted.bam'
      ]
    if clean_r2:
        fc_cmd += ['-p', '--countReadPairs']

    fc_cmd += ['alignments.sorted.bam']

    try:
        run(fc_cmd, capture_output=True)
    except subprocess.CalledProcessError as e:
        print("---- STDOUT ----")
        print(e.stdout)
        print("---- STDERR ----")
        print(e.stderr)
        raise

    df_counts, total_counted = load_featurecounts(fc_out)
    print(f"Counts Table list. Total counted = {total_counted}")
    cp_to_checkpoint(fc_out, 'gene_counts.txt')

# ==============================================================================
# 13) PARSE GFF
# ==============================================================================

print("\nParse GFF...")
df_meta = parse_annotation_gff(gff_file, fc_group_attr)
if df_meta.empty:
    raise RuntimeError("CDS not extracted from GFF.")

# # ==============================================================================
# 14) KOFAMSCAN -> KO
# ==============================================================================

print("\nRunning KOfamScan...")
kofam_result = 'kofam_raw.tsv'
kofam_profiles_dir = os.path.join(KOFAM_DB_DIR, 'profiles')
kofam_ko_list = os.path.join(KOFAM_DB_DIR, 'ko_list')

if not FORCE_RERUN and checkpoint_exists('kofam_raw.tsv'):
    restore_from_checkpoint('kofam_raw.tsv', kofam_result)
    df_kofam = parse_kofam_output(kofam_result)
    print(f"Restoring KOfam from checkpoint. Genes with KO: {len(df_kofam)}")
else:
    kofam_ok = os.path.exists(kofam_profiles_dir) and os.path.exists(kofam_ko_list)

    if not kofam_ok:
        print("KOfam DB not disponible.")
        df_kofam = pd.DataFrame(columns=['gene', 'ko', 'ko_desc'])
    else:
        cmd = [
            os.path.join(KOFAM_SCAN_DIR, 'exec_annotation'),
            '--cpu', '4',
            '--profile', kofam_profiles_dir,
            '--ko-list', kofam_ko_list,
            '-o', kofam_result,
            faa_file
        ]
        run(cmd)
        print("KOfamScan completed.")

        df_kofam = parse_kofam_output(kofam_result)
        print(f"Genes with KO: {len(df_kofam)}")
        cp_to_checkpoint(
    kofam_result,
    'kofam_raw.tsv'
)
print("\n===== AUDITORÍA KOFAM =====")

suspect_terms = [
    'adrenergic',
    'laminin',
    'myosin',
    'integrin',
    'cadherin',
    'apolipoprotein'
]

for term in suspect_terms:

    tmp = df_kofam[
        df_kofam['ko_desc']
        .astype(str)
        .str.lower()
        .str.contains(term, na=False)
    ]

    if len(tmp) > 0:

        print(f"\nPosible hits: {term}")
        print(
            tmp[['gene','ko','ko_desc']]
            .head(20)
            .to_string(index=False)
        )
        cp_to_checkpoint(kofam_result, 'kofam_raw.tsv')

# Normalize genes from KOfam
if not df_kofam.empty:
    df_kofam['gene'] = df_kofam['gene'].apply(normalize_gene_id)


# ==============================================================================
# 15) KEGG: KO -> PATHWAY (rápido + cache local)
# ==============================================================================
print("\nRecovering KO maps -> from KEGG pathways...")

pathway_name_map = {}
ko_to_pathways = {}

if not df_kofam.empty:

    try:

        pathway_name_map = fetch_kegg_pathway_names()

        print("\n===== PATHWAYS PROBLEMÁTICOS =====")

        for pid in [
            'ko04913',
            'ko05204',
            'ko05207',
            'ko05208',
            'ko05417'
        ]:

            print(
                pid,
                pathway_name_map.get(pid, 'NO FOUND')
            )

        kos_presentes = sorted(
            set([
                k for k in df_kofam['ko']
                if pd.notna(k)
            ])
        )

        ko_to_pathways = fetch_kegg_pathways_for_kos_fast(
            kos_presentes
        )

        print("Map KEGG ready.")
        print(
            f"KOs unique at sample: "
            f"{len(set(kos_presentes))}"
        )

        print(
            f"KOs at least one pathway: "
            f"{sum(1 for k, v in ko_to_pathways.items() if v)}"
        )

    except Exception as e:

        print(
            f"KEGG fail pathway mapping: {e}"
        )

        pathway_name_map = {}
        ko_to_pathways = {}

else:

    print(
        "df_kofam is empty; "
        "no KO to map to KEGG."
    )

# ==============================================================================
# 15B) FINAL VALIDATION FAA/GFF Y KOFAM
# ==============================================================================

# Normalizar genes actuales del GFF/meta y gene_counts
df_meta['gene'] = df_meta['gene'].apply(normalize_gene_id)
df_counts['gene'] = df_counts['gene'].apply(normalize_gene_id)

current_genes = set(df_meta['gene'].astype(str))
faa_genes = get_faa_gene_ids(faa_file)
kofam_genes = set(df_kofam['gene'].astype(str)) if not df_kofam.empty else set()

overlap_gff_kofam = current_genes.intersection(kofam_genes)
overlap_faa_kofam = faa_genes.intersection(kofam_genes)

print(f"Genes actuales en GFF/meta: {len(current_genes)}")
print(f"Genes en FAA actual: {len(faa_genes)}")
print(f"Genes en KOfam: {len(kofam_genes)}")
print(f"Overlap GFF/KOfam: {len(overlap_gff_kofam)}")
print(f"Overlap FAA/KOfam: {len(overlap_faa_kofam)}")

print("\nPrimeros genes de df_meta:")
print(df_meta['gene'].astype(str).head(10).to_list())

if not df_kofam.empty:
    print("\nPrimeros genes de df_kofam:")
    print(df_kofam['gene'].astype(str).head(10).to_list())

# Si KOfam no coincide con el universo actual de genes, descartar
if len(current_genes) > 0:
    overlap_fraction = len(overlap_gff_kofam) / len(current_genes)
else:
    overlap_fraction = 0.0

if len(overlap_gff_kofam) == 0 or overlap_fraction < 0.10:
    print("NO match KOfam with current genes.")
    print("df_kofam will be deleted to aoid errors.")
    df_kofam = pd.DataFrame(columns=['gene', 'ko', 'ko_desc'])
    ko_to_pathways = {}
    pathway_name_map = {}


# ==============================================================================
# 16) TABLE INTEGRATION
# ==============================================================================

def suspicious_ko(desc):

    desc = str(desc).lower()

    blacklist = [
        'adrenergic receptor',
        'laminin',
        'myosin',
        'cadherin',
        'integrin',
        'apolipoprotein',
        'nuclear antigen',
        'olfactory receptor',
        'neuroactive ligand',
        'vascular smooth muscle',
        'cardiomyocyte'
        "laminin",
        "integrin",
        "dynein",
        "glutamate receptor",
        "nmda",
        "anoctamin",
        "glucagon",
        "galanin receptor",
        "adhesion g-protein",
        "g-protein coupled receptor",
        "intercellular adhesion molecule",
        "kinetochore",
        "tryparedoxin",
        "syntaxin",
        "fibrillin",
        "endoplasmic reticulum",
        "heterogeneous nuclear",
        "receptor-type tyrosine"
    ]

    return any(x in desc for x in blacklist)
# Suspected hits
if not df_kofam.empty:

    df_kofam['ko_suspicious'] = (
        df_kofam['ko_desc']
        .apply(suspicious_ko)
    )

    print(
        f"Detected suspicious KO: "
        f"{df_kofam['ko_suspicious'].sum()}"
    )

    suspect_kos = df_kofam[
        df_kofam['ko_suspicious']
    ].copy()

    suspect_kos.to_csv(
        'KOFAM_SUSPECT_HITS.csv',
        index=False
    )

else:
    df_kofam['ko_suspicious'] = False


print("\nIntegrando conteos, metadatos y KO/pathways...")

df = df_counts.merge(df_meta, on='gene', how='left')

if not df_kofam.empty:
    df = df.merge(df_kofam, on='gene', how='left')
else:
    df['ko'] = pd.NA
    df['ko_desc'] = pd.NA
    df['ko_suspicious'] = False


def map_gene_pathways(ko):

    if pd.isna(ko):
        return []

    return ko_to_pathways.get(ko, [])


# ==============================================================================
# Taking out no bacteria pathways
# ==============================================================================

BAD_PREFIXES = (

    # Sistemas multicelulares
    'ko045', 'ko046', 'ko047', 'ko048', 'ko049',
    'map045', 'map046', 'map047', 'map048', 'map049',

    # Enfermedades humanas
    'ko050', 'ko051', 'ko052', 'ko053', 'ko054',
    'map050', 'map051', 'map052', 'map053', 'map054'
)


def keep_bacterial_pathways(pathways):

    if not isinstance(pathways, list):
        return []

    return [
        pw
        for pw in pathways
        if not str(pw).startswith(BAD_PREFIXES)
    ]


# ===============================
# mapping pathways
# ===============================

df['pathways'] = df['ko'].apply(map_gene_pathways)


# ===============================
# Fintering eukariotik pathways
# ===============================

df['pathways'] = df['pathways'].apply(
    keep_bacterial_pathways
)


# ===============================
# Desactivar pathways de KO suspicious
# ===============================

mask = (df['ko_suspicious'] == True)

for idx in df.index[mask]:
    df.at[idx, 'pathways'] = []


# ===============================
# recalculate pathways
# ===============================

df['pathway_count'] = df['pathways'].apply(len)

print("\n=== DEBUG INTEGRACIÓN KO/PATHWAYS ===")
print("Columnas en df:", list(df.columns))
print("Genes totales en df:", len(df))
print("Genes con KO no nulo:", df['ko'].notna().sum() if 'ko' in df.columns else 'NA')
print("Genes con pathway_count > 0:", (df['pathway_count'] > 0).sum() if 'pathway_count' in df.columns else 'NA')

print("\nEjemplos con KO:")

if 'ko' in df.columns:
    print(
        df.loc[
            df['ko'].notna(),
            ['gene', 'gene_name', 'product', 'ko', 'ko_desc', 'ko_suspicious']
        ]
        .head(20)
        .to_string(index=False)
    )


print("\nEjemplos con pathways:")

if 'pathways' in df.columns:

    tmp_pw = df.loc[
        df['pathway_count'] > 0,
        [
            'gene',
            'gene_name',
            'product',
            'ko',
            'ko_suspicious',
            'pathways',
            'pathway_count'
        ]
    ].head(20)

    print(tmp_pw.to_string(index=False))


print("\nEjemplos sin pathways pero con KO:")

if 'pathways' in df.columns and 'ko' in df.columns:

    tmp_miss = df.loc[
        (df['ko'].notna()) &
        (df['pathway_count'] == 0),
        [
            'gene',
            'gene_name',
            'product',
            'ko',
            'ko_desc',
            'ko_suspicious'
        ]
    ].head(20)

    print(tmp_miss.to_string(index=False))

# ==============================================================================
# 17) Heuristic cat
# ==============================================================================

print("\nAsignando categorías heurísticas...")
df['cat'] = df.apply(classify_gene, axis=1)

# ==============================================================================
# 18) EXTERNAL CORE-LIKE: MATCHING reference
# ==============================================================================

if external_core_ref is not None:
    print(f"\nCargando referencia externa: {external_core_ref}")
    df_ref = load_core_reference(external_core_ref)

    match_results = df.apply(lambda row: match_against_core_reference(row, df_ref), axis=1)
    match_df = pd.DataFrame(match_results.tolist(), columns=[
        'external_core_match',
        'external_core_match_type',
        'external_core_ref_locus',
        'external_core_ref_gene',
        'external_core_ref_description'
    ])

    df = pd.concat([df.reset_index(drop=True), match_df.reset_index(drop=True)], axis=1)
else:
    df['external_core_match'] = False
    df['external_core_match_type'] = ""
    df['external_core_ref_locus'] = ""
    df['external_core_ref_gene'] = ""
    df['external_core_ref_description'] = ""
# ==============================================================================
# 18B) MODE CL - CH  CORE EXTERNO and Abundance
# ==============================================================================

df['external_core_expr_band'] = ""

if 'external_core_match' in df.columns and 'RPKM' in df.columns:
    core_ext = df[df['external_core_match'] == True].copy()

    if not core_ext.empty:
        q25 = core_ext['RPKM'].quantile(0.25)
        q75 = core_ext['RPKM'].quantile(0.75)

        df.loc[
            (df['external_core_match'] == True) & (df['RPKM'] <= q25),
            'external_core_expr_band'
        ] = 'low'

        df.loc[
            (df['external_core_match'] == True) & (df['RPKM'] >= q75),
            'external_core_expr_band'
        ] = 'high'

        df.loc[
            (df['external_core_match'] == True) &
            (df['RPKM'] > q25) &
            (df['RPKM'] < q75),
            'external_core_expr_band'
        ] = 'mid'

        print(f"Core externo q25 RPKM: {q25:.4f}")
        print(f"Core externo q75 RPKM: {q75:.4f}")
        print("Distribución de bandas de expresión en core externo:")
        print(df.loc[df['external_core_match'] == True, 'external_core_expr_band'].value_counts(dropna=False))
    else:
        print("No external genes with external_core_match=True; no se calcularon bandas de expresión.")
else:
    print("No se pudo calcular external_core_expr_band (faltan columnas external_core_match o RPKM).")
# ==============================================================================
# 19) SELECT-NET filters
# ==============================================================================
print("\nAppying finltering mode...")

if mode in ['C', 'CH', 'CL']:
    if 'core_like' not in df.columns:
        df['core_like'] = False

    df['core_like'] = df['core_like'] | df['external_core_match']

def filter_select_net(df, mode):
    if mode == 'A':
        return df[df['cat'].isin([
            'Extremophile-like',
            'Bioplastic-like',
            'Biodegradation-like',
            'Industrial-enzyme-like'
        ])].copy()

    elif mode == 'B':
        q = df['RPKM'].quantile(0.90)
        return df[df['RPKM'] >= q].copy()

    elif mode == 'C':
        return df[~df['core_like']].copy()

    elif mode == 'CH':
        return df[~(
            (df['external_core_match'] == True) &
            (df['external_core_expr_band'].isin(['mid', 'high']))
        )].copy()

    elif mode == 'CL':
        return df[~(
            (df['external_core_match'] == True) &
            (df['external_core_expr_band'].isin(['low', 'mid']))
        )].copy()

    elif mode == 'D':
        return df[df['cat'].isin(['Resistance', 'Virulence'])].copy()

    else:
        return df.copy()

# ------------------------------------------------------------------------------
# APPLY FILTER
# ------------------------------------------------------------------------------
df_filtered = filter_select_net(df, mode)

if df_filtered.empty:
    print(f"The mode {mode} remove all genes. Full network will be provided.")
    df_filtered = df.copy()

# Save Seeds (MUY IMPORTANTE)
original_df_filtered = df_filtered.copy()
seed_gene_count = len(original_df_filtered)
# ------------------------------------------------------------------------------
# Recover neightborhoods (SIMPLE)
# ------------------------------------------------------------------------------
extra_genes = set()

for _, row in df_filtered.iterrows():
    contig = row['contig']
    start = row['start']

    vecinos = df[
        (df['contig'] == contig) &
        (df['start'] >= start - 2000) &
        (df['start'] <= start + 2000)
    ]

    for g in vecinos['gene']:
        extra_genes.add(g)

# combinar semillas + vecinos
df_filtered = df[df['gene'].isin(set(df_filtered['gene']).union(extra_genes))]

print(f"Genes tras recuperar vecinos: {len(df_filtered)}")
context_expanded_gene_count = len(df_filtered)

# ==============================================================================
# 20) Genomic Proximity Based-Network
# ==============================================================================
print("\nBuilding Genomic Proximity Based-Network")

import networkx as nx

G = nx.Graph()

# nodos
all_genes = df_filtered['gene'].dropna().unique().tolist()
G.add_nodes_from(all_genes)

# atributos
for _, row in df_filtered.iterrows():
    gene = row['gene']
    if gene in G:
        G.nodes[gene]['category'] = row.get('cat', 'Other')
        G.nodes[gene]['rpkm'] = row.get('RPKM', 0)
        G.nodes[gene]['gene_name'] = row.get('gene_name', '')

# ordenar
df_pos = df_filtered.dropna(subset=['contig', 'start', 'end']).copy()
df_pos = df_pos.sort_values(['contig', 'start'])

distance_threshold = 2000
edges_added = 0

# edges por proximidad
for contig, group in df_pos.groupby('contig'):
    group = group.sort_values('start').reset_index(drop=True)

    for i in range(len(group) - 1):
        g1 = group.iloc[i]
        g2 = group.iloc[i + 1]

        distance = g2['start'] - g1['end']

        if distance < distance_threshold:
            G.add_edge(
                g1['gene'],
                g2['gene'],
                weight=1,
                edge_type="genomic_proximity"
            )
            edges_added += 1

print(f"Edges por proximidad añadidos: {edges_added}")
print(f"Nodos totales: {G.number_of_nodes()}")
print(f"Edges iniciales: {G.number_of_edges()}")


# ------------------------------------------------------------------------------
# EDGES filter
# ------------------------------------------------------------------------------
def get_prefix(name, n=3):
    if pd.isna(name):
        return None
    return str(name).lower()[:n]

edges_to_remove = []

for u, v in G.edges():
    name_u = G.nodes[u].get('gene_name')
    name_v = G.nodes[v].get('gene_name')

    prefix_u = get_prefix(name_u)
    prefix_v = get_prefix(name_v)

    if prefix_u and prefix_v and prefix_u != prefix_v:
        edges_to_remove.append((u, v))

G.remove_edges_from(edges_to_remove)

print(f"Edges eliminados por prefijo: {len(edges_to_remove)}")
print(f"Edges tras limpieza: {G.number_of_edges()}")


# ------------------------------------------------------------------------------
# Seed filters (CLAVE FINAL)
# ------------------------------------------------------------------------------
seed_genes = set(original_df_filtered['gene'])

components_to_remove = []

for comp in nx.connected_components(G):
    if not any(g in seed_genes for g in comp):
        components_to_remove.extend(comp)

G.remove_nodes_from(components_to_remove)

print(f"Componentes sin semilla eliminados: {len(components_to_remove)}")
print(f"Nodos finales: {G.number_of_nodes()}")
print(f"Edges finales: {G.number_of_edges()}")

# ==============================================================================
# 21) EXPORT TABLES
# =============================================================================
# Añadir etiquetas amigables
full_csv = os.path.join(
    OUT_DIR,
    f"{sample_id}_SELECTNET_full_table.csv"
)

filt_csv = os.path.join(
    OUT_DIR,
    f"{sample_id}_SELECTNET_mode_{mode}_table.csv"
)

df.to_csv(
    full_csv,
    index=False
)

df_filtered.to_csv(
    filt_csv,
    index=False
)

print(f"Tabla completa exportada a: {full_csv}")
print(f"Tabla filtrada exportada a: {filt_csv}")

# ==============================================================================
# 21A) EXPORTAR RED COMO NODOS / ARISTAS (ENRIQUECIDOS)
# ==============================================================================

print("\nExportando red como archivos reutilizables (enriquecidos)...")

# ----------------------------------------------------------------------
# 1) Construir mapas robustos de anotación por gen
# ----------------------------------------------------------------------

def safe_list(x):
    if isinstance(x, list):
        return x
    if pd.isna(x):
        return []
    if isinstance(x, str):
        if ';' in x:
            return [i.strip() for i in x.split(';') if i.strip()]
        elif x.strip():
            return [x.strip()]
    return []

# Mapa gene -> atributos desde df_filtered
node_attr_map = {}

for _, row in df_filtered.iterrows():
    gene = row['gene']
    pathways_list = safe_list(row.get('pathways', []))

    node_attr_map[gene] = {
        'rpkm': row.get('RPKM', None),
        'count': row.get('count', None),
        'category': row.get('cat', row.get('category', None)),
        'product': row.get('product', None),
        'gene_name': row.get('gene_name', None),
        'ko': row.get('ko', None),
        'ko_desc': row.get('ko_desc', None),
        'pathways': pathways_list,
        'pathway_names': [pathway_name_map.get(p, '') for p in pathways_list],
        'core_like': row.get('core_like', None)
    }

# ----------------------------------------------------------------------
# 2) EXPORTAR NODOS ENRIQUECIDOS
# ----------------------------------------------------------------------

# ==============================================================================
# NODE LABELS TO CYTOSCAPE
# ==============================================================================

label_counters = {}

def next_label(prefix):

    if prefix not in label_counters:
        label_counters[prefix] = 1
    else:
        label_counters[prefix] += 1

    return f"{prefix}{label_counters[prefix]}"


def assign_node_label(row):
    ko_desc = str(row.get('ko_desc', '')).strip()
    gene = str(row.get('gene', '')).strip()
    gene_name = str(row.get('gene_name', '')).strip()
    product = str(row.get('product', '')).strip()
    product_lower = product.lower()

    # --------------------------------------------------
    # IDs long ids from NCBI
    # --------------------------------------------------

    is_ncbi_id = (
        gene_name.startswith("cds-WP_")
        or gene_name.startswith("WP_")
        or gene.startswith("WP_")
    )

    # --------------------------------------------------
    # PRIORIDAD 1
    # Nombre de gen real
    # --------------------------------------------------

    if (
        gene_name
        and gene_name != gene
        and not is_ncbi_id
    ):
        return gene_name
# --------------------------------------------------
    # PRIORIDAD 2
    # Buscar nombre corto útil en PRODUCT
    # --------------------------------------------------

    if (
        product
        and "hypothetical protein" not in product_lower
    ):

        words = (
            product
            .replace('-', ' ')
            .replace('/', ' ')
            .split()
        )

        BAD_LABELS = {
            "ABC",
            "ATP",
            "kDa",
            "protein",
            "family",
            "binding"
        }
        for w in reversed(words):
            w_clean = w.strip("(),.;:[]")

            if (
                len(w_clean) >= 3
                and any(c.isupper() for c in w_clean)
                and w_clean not in BAD_LABELS
                and "kDa" not in w_clean
            ):
                return w_clean

    # --------------------------------------------------
    # PRIORIDAD 3
    # EC NUMBER
    # --------------------------------------------------

    if "[EC:" in ko_desc:

        start = ko_desc.find("[EC:")
        end = ko_desc.find("]", start)

        if end > start:
            return ko_desc[start+1:end]

    # --------------------------------------------------
    # PRIORIDAD 4
    # KO
    # --------------------------------------------------

    ko = str(row.get('ko', '')).strip()

    if ko.startswith("K"):
        return ko
    # --------------------------------------------------
    # PRIORIDAD 4
    # Etiquetas funcionales heredadas
    # --------------------------------------------------

    if 'type iii secretion' in product_lower:
        return next_label('T3S')

    if 'type vi secretion' in product_lower:
        return next_label('T6S')

    if 'fimbr' in product_lower:
        return next_label('FIM')

    if 'siderophore' in product_lower:
        return next_label('SID')

    if 'efflux' in product_lower:
        return next_label('MDR')

    if 'adhesin' in product_lower:
        return next_label('ADH')

    # --------------------------------------------------
    # PRIORIDAD 5
    # Hipotéticas
    # --------------------------------------------------

    return next_label('HP')


nodes_rows = []

for n, data in G.nodes(data=True):

    anno = node_attr_map.get(n, {})

    pathways_list = anno.get('pathways', data.get('pathways', []))

    if not isinstance(pathways_list, list):
        pathways_list = safe_list(pathways_list)

    pathway_names_list = anno.get('pathway_names', [])

    if not isinstance(pathway_names_list, list):
        pathway_names_list = safe_list(pathway_names_list)

    ko_val = anno.get('ko', data.get('ko', None))

    if isinstance(ko_val, list):
        ko_val = ';'.join(ko_val)

    node_row = {
        'gene': n,
        'rpkm': anno.get('rpkm', data.get('rpkm', None)),
        'count': anno.get('count', data.get('count', None)),
        'category': anno.get('category', data.get('category', None)),
        'product': anno.get('product', data.get('product', None)),
        'gene_name': anno.get('gene_name', None),
        'ko': ko_val,
        'ko_desc': anno.get('ko_desc', None),
        'pathways': ';'.join(pathways_list) if pathways_list else '',
        'pathway_names': ';'.join([p for p in pathway_names_list if p]) if pathway_names_list else '',
        'pathway_count': len(pathways_list),
        'core_like': anno.get('core_like', data.get('core_like', None)),
        'degree_in_exported_graph': G.degree(n)
    }

    nodes_rows.append(node_row)

df_nodes = pd.DataFrame(nodes_rows)

# firendly labels
df_nodes['node_label'] = df_nodes.apply(
    assign_node_label,
    axis=1
)
label_map = dict(
    zip(
        df_nodes['gene'],
        df_nodes['node_label']
    )
)
df['node_label'] = (
    df['gene']
    .map(label_map)
)

df_filtered['node_label'] = (
    df_filtered['gene']
    .map(label_map)
)

nodes_csv = os.path.join(
    OUT_DIR,
    f"{sample_id}_SELECTNET_mode_{mode}_nodes.csv"
)

df_nodes.to_csv(
    nodes_csv,
    index=False
)
# ----------------------------------------------------------------------
# 3) EXPORTAR ARISTAS ENRIQUECIDAS
# ----------------------------------------------------------------------

def infer_edge_type(source_gene, target_gene, shared_pathways):
    """
    NO crea links nuevos.
    Solo clasifica mejor los links ya existentes.
    """
    s = node_attr_map.get(source_gene, {})
    t = node_attr_map.get(target_gene, {})

    s_cat = s.get('category', None)
    t_cat = t.get('category', None)

    s_ko = s.get('ko', None)
    t_ko = t.get('ko', None)

    if shared_pathways:
        return 'shared_pathway'
    if pd.notna(s_ko) and pd.notna(t_ko) and str(s_ko).strip() and str(t_ko).strip() and s_ko == t_ko:
        return 'same_KO'
    if s_cat and t_cat and s_cat == t_cat and s_cat != 'Other':
        return 'same_category_fallback'
    return 'existing_link_unresolved'

edges_rows = []
for u, v, data in G.edges(data=True):
    u_anno = node_attr_map.get(u, {})
    v_anno = node_attr_map.get(v, {})

    u_paths = u_anno.get('pathways', [])
    v_paths = v_anno.get('pathways', [])

    if not isinstance(u_paths, list):
        u_paths = safe_list(u_paths)
    if not isinstance(v_paths, list):
        v_paths = safe_list(v_paths)

    # recalcular shared pathways por si el edge original venía sin la lista
    recalculated_shared = sorted(set(u_paths).intersection(set(v_paths)))

    shared_raw = data.get('shared_pathways', [])
    if not isinstance(shared_raw, list):
        shared_raw = safe_list(shared_raw)

    final_shared = sorted(set(shared_raw).union(set(recalculated_shared)))

    u_ko = u_anno.get('ko', None)
    v_ko = v_anno.get('ko', None)

    if isinstance(u_ko, list):
        u_ko = ';'.join(u_ko)
    if isinstance(v_ko, list):
        v_ko = ';'.join(v_ko)

    edge_type = infer_edge_type(u, v, final_shared)

    edges_rows.append({
        'source': u,
        'target': v,
        'weight': data.get('weight', 1),

        # info original / recomputada
        'shared_pathways': ';'.join(final_shared) if final_shared else '',
        'shared_pathway_names': ';'.join([pathway_name_map.get(p, '') for p in final_shared if p]) if final_shared else '',
        'shared_pathway_count': len(final_shared),

        # atributos del nodo source
        'source_category': u_anno.get('category', None),
        'source_product': u_anno.get('product', None),
        'source_gene_name': u_anno.get('gene_name', None),
        'source_ko': u_ko,
        'source_pathways': ';'.join(u_paths) if u_paths else '',

        # atributos del nodo target
        'target_category': v_anno.get('category', None),
        'target_product': v_anno.get('product', None),
        'target_gene_name': v_anno.get('gene_name', None),
        'target_ko': v_ko,
        'target_pathways': ';'.join(v_paths) if v_paths else '',

        # interpretación mínima del edge ya existente
        'edge_type_inferred': edge_type
    })

df_edges = pd.DataFrame(edges_rows)

# ----------------------------------------------------------------------
# safe DEBUG: after df_edges
# ----------------------------------------------------------------------
print("\n[DEBUG] Columnas en df_edges:")
print(df_edges.columns.tolist())

print("\n[DEBUG] Primeras filas de df_edges:")
print(df_edges.head())

print("\n[DEBUG] Tamaño de df_edges:")
print(df_edges.shape)

# asegurar columnas mínimas por si la red quedó vacía
if df_edges.empty:
    df_edges = pd.DataFrame(columns=[
        'source', 'target', 'weight',
        'shared_pathways', 'shared_pathway_names', 'shared_pathway_count',
        'source_category', 'source_product', 'source_gene_name', 'source_ko', 'source_pathways',
        'target_category', 'target_product', 'target_gene_name', 'target_ko', 'target_pathways',
        'edge_type_inferred'
    ])

edges_csv = os.path.join(OUT_DIR, f"{sample_id}_SELECTNET_mode_{mode}_edges.csv")
df_edges.to_csv(edges_csv, index=False)

print(f"Nodos enriquecidos exportados a: {nodes_csv}")
print(f"Aristas enriquecidas exportadas a: {edges_csv}")

# ==============================================================================
# 21B) EXPORTAR ARISTAS GLOBALES (SIN FILTRO) BASADAS EN PATHWAYS COMPARTIDOS
# ==============================================================================

print("\nExportando aristas globales (sin filtro) basadas en pathways compartidos...")

# ---------------------------------------------------------------------
# 1) Seleccionar universo global para red
#    Recomendado: solo genes con pathway asignado
# ---------------------------------------------------------------------
df_global_net = df[df['pathway_count'] > 0].copy()

# Si alguna vez quisieras incluir todos los genes, incluso sin pathway:
# df_global_net = df.copy()

print(f"Genes totales en df: {len(df)}")
print(f"Genes incluidos en la red global (pathway_count > 0): {len(df_global_net)}")

# ---------------------------------------------------------------------
# 2) Construir red global SOLO por pathways compartidos
#    (sin fallback por categoría)
# ---------------------------------------------------------------------
G_global = build_kegg_gene_graph(df_global_net)

print(f"Red global construida: Nodes={G_global.number_of_nodes()} | Edges={G_global.number_of_edges()}")

# ---------------------------------------------------------------------
# 3) Helper local por si pathways vienen como lista o string
# ---------------------------------------------------------------------
def safe_list_global(x):
    if isinstance(x, list):
        return x
    if pd.isna(x):
        return []
    if isinstance(x, str):
        if ';' in x:
            return [i.strip() for i in x.split(';') if i.strip()]
        elif x.strip():
            return [x.strip()]
    return []

# ---------------------------------------------------------------------
# 4) Mapa global gene -> atributos
# ---------------------------------------------------------------------
global_attr_map = {}
for _, row in df_global_net.iterrows():
    gene = row['gene']
    pathways_list = safe_list_global(row.get('pathways', []))
    global_attr_map[gene] = {
        'rpkm': row.get('RPKM', None),
        'count': row.get('count', None),
        'category': row.get('cat', row.get('category', None)),
        'product': row.get('product', None),
        'gene_name': row.get('gene_name', None),
        'ko': row.get('ko', None),
        'ko_desc': row.get('ko_desc', None),
        'pathways': pathways_list,
        'pathway_names': [pathway_name_map.get(p, '') for p in pathways_list],
        'core_like': row.get('core_like', None)
    }

# ---------------------------------------------------------------------
# 5) Exportar just aristas globales
# ---------------------------------------------------------------------
global_edges_rows = []

for u, v, data in G_global.edges(data=True):
    u_anno = global_attr_map.get(u, {})
    v_anno = global_attr_map.get(v, {})

    u_paths = u_anno.get('pathways', [])
    v_paths = v_anno.get('pathways', [])

    if not isinstance(u_paths, list):
        u_paths = safe_list_global(u_paths)
    if not isinstance(v_paths, list):
        v_paths = safe_list_global(v_paths)

    recalculated_shared = sorted(set(u_paths).intersection(set(v_paths)))

    shared_raw = data.get('shared_pathways', [])
    if not isinstance(shared_raw, list):
        shared_raw = safe_list_global(shared_raw)

    final_shared = sorted(set(shared_raw).union(set(recalculated_shared)))

    u_ko = u_anno.get('ko', None)
    v_ko = v_anno.get('ko', None)

    if isinstance(u_ko, list):
        u_ko = ';'.join(u_ko)
    if isinstance(v_ko, list):
        v_ko = ';'.join(v_ko)

    global_edges_rows.append({
        'source': u,
        'target': v,
        'weight': data.get('weight', 1),
        'shared_pathways': ';'.join(final_shared) if final_shared else '',
        'shared_pathway_names': ';'.join([pathway_name_map.get(p, '') for p in final_shared if p]) if final_shared else '',
        'shared_pathway_count': len(final_shared),

        'source_category': u_anno.get('category', None),
        'source_product': u_anno.get('product', None),
        'source_gene_name': u_anno.get('gene_name', None),
        'source_ko': u_ko,
        'source_pathways': ';'.join(u_paths) if u_paths else '',

        'target_category': v_anno.get('category', None),
        'target_product': v_anno.get('product', None),
        'target_gene_name': v_anno.get('gene_name', None),
        'target_ko': v_ko,
        'target_pathways': ';'.join(v_paths) if v_paths else '',

        'edge_type_inferred': 'shared_pathway_global'
    })

df_global_edges = pd.DataFrame(global_edges_rows)

if df_global_edges.empty:
    df_global_edges = pd.DataFrame(columns=[
        'source', 'target', 'weight',
        'shared_pathways', 'shared_pathway_names', 'shared_pathway_count',
        'source_category', 'source_product', 'source_gene_name', 'source_ko', 'source_pathways',
        'target_category', 'target_product', 'target_gene_name', 'target_ko', 'target_pathways',
        'edge_type_inferred'
    ])

global_edges_csv = os.path.join(OUT_DIR, f"{sample_id}_SELECTNET_global_edges.csv")
df_global_edges.to_csv(global_edges_csv, index=False)

print(f"Aristas globales exportadas a: {global_edges_csv}")
print(f"Total de aristas globales exportadas: {len(df_global_edges)}")

# ----------------------------------------------------------------------
# 4) EXPORT DE GRAPHML / GEXF con atributos enriquecidos
# ----------------------------------------------------------------------

G_export = nx.Graph()

for _, row in df_nodes.iterrows():
    gene = row['gene']
    attrs = row.to_dict()

    clean_attrs = {}
    for k, v in attrs.items():
        if pd.isna(v):
            clean_attrs[k] = ''
        else:
            clean_attrs[k] = v
    G_export.add_node(gene, **clean_attrs)

for _, row in df_edges.iterrows():
    src = row['source']
    tgt = row['target']
    attrs = row.to_dict()

    clean_attrs = {}
    for k, v in attrs.items():
        if pd.isna(v):
            clean_attrs[k] = ''
        else:
            clean_attrs[k] = v
    G_export.add_edge(src, tgt, **clean_attrs)

graphml_out = os.path.join(OUT_DIR, f"{sample_id}_SELECTNET_mode_{mode}.graphml")
nx.write_graphml(G_export, graphml_out)

gexf_out = os.path.join(OUT_DIR, f"{sample_id}_SELECTNET_mode_{mode}.gexf")
nx.write_gexf(G_export, gexf_out)

print(f"GraphML enriquecido exportado a: {graphml_out}")
print(f"GEXF enriquecido exportado a: {gexf_out}")

# ----------------------------------------------------------------------
# 5) TABLA gene -> KO -> pathway
# ----------------------------------------------------------------------

gene_ko_path_rows = []
for _, row in df_filtered.iterrows():
    pws = row['pathways'] if isinstance(row['pathways'], list) else []
    if not pws:
        gene_ko_path_rows.append({
            'gene': row['gene'],
            'ko': row.get('ko', ''),
            'pathway_id': '',
            'pathway_name': ''
        })
    else:
        for pw in pws:
            gene_ko_path_rows.append({
                'gene': row['gene'],
                'ko': row.get('ko', ''),
                'pathway_id': pw,
                'pathway_name': pathway_name_map.get(pw, '')
            })

df_gene_ko_path = pd.DataFrame(gene_ko_path_rows)
gene_ko_path_csv = os.path.join(OUT_DIR, f"{sample_id}_SELECTNET_mode_{mode}_gene_ko_pathway.csv")
df_gene_ko_path.to_csv(gene_ko_path_csv, index=False)

print(f"Tabla gene->KO->pathway exportada a: {gene_ko_path_csv}")

# ==============================================================================
# 22) VISUALIZACIÓN - just to see as an example of the process graphxlm to cytoscape
# ==============================================================================

print("\nDibujando red...")

plt.figure(figsize=(14, 12))
pos = nx.kamada_kawai_layout(G)

top_nodes = sorted(G.nodes(), key=lambda n: G.nodes[n].get('rpkm', 0), reverse=True)[:20]
labels = {n: G.nodes[n].get('gene_name') or n for n in top_nodes}

color_map = {
    'Extremophile-like': '#FF5733',
    'Bioplastic-like': '#2ECC71',
    'Biodegradation-like': '#3498DB',
    'Industrial-enzyme-like': '#F39C12',
    'Resistance': '#FFC300',
    'Virulence': '#8E44AD',
    'Other': '#5DADE2'
}

node_colors = [color_map.get(G.nodes[n]['category'], 'gray') for n in G.nodes()]
node_sizes = [max(80, G.nodes[n]['rpkm'] * 20) for n in G.nodes()]
edge_widths = [max(0.5, d.get('weight', 1)) for _, _, d in G.edges(data=True)]

nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=node_sizes, alpha=0.85)
nx.draw_networkx_edges(G, pos, width=edge_widths, alpha=0.35)
nx.draw_networkx_labels(G, pos, labels=labels, font_size=7)

interpretation_label = "Expression-aware" if lib_mode == 'R' else "Abundance-aware"

plt.title(
    f"SELECT-NET | {sample_id} | Mode {mode} | {interpretation_label}\n"
    f"Nodes={G.number_of_nodes()} | Edges={G.number_of_edges()} | External-core={external_core_ref}"
)

plt.axis('off')
plt.tight_layout()

png_out = os.path.join(OUT_DIR, f"{sample_id}_SELECTNET_mode_{mode}.png")
plt.savefig(png_out, dpi=300, bbox_inches='tight')
plt.show()

print(f"Red DONE: {png_out}")

# ==============================================================================
# 23) Summary
# ==============================================================================

print("\nPIPELINE COMPLETED!!!!!!!!")
print(f"Sample: {sample_id}")
print(f"Mode: {mode}")
print(f"External core reference: {external_core_ref}")
print(f"External core ref applied to filter: {external_core_ref_applied_to_filter}")
print(f"Annotation mode: {annotation_mode}")
print(f"Total Genes: {len(df)}")
print(f"Filteres Genes: {len(df_filtered)}")
print(f"Nodes: {G.number_of_nodes()}")
print(f"Edges: {G.number_of_edges()}")
print(f"Interpretation: {'Expresión (RNA-seq)' if lib_mode == 'R' else 'Abundance normalized (WGS/metagenoma)'}")
print(f"Checkpoints at Drive: {CHECKPOINT_DIR}")

if mode in ['C', 'CH', 'CL'] and 'external_core_expr_band' in df.columns:
    print("\nDistribución de bandas de expresión en core externo:")
    print(df.loc[df['external_core_match'] == True, 'external_core_expr_band'].value_counts(dropna=False))

# ==============================================================================
# 24) Runing report
# ==============================================================================

import json
from datetime import datetime

print("\n Runing report in progress")

core_matches_csv = ""

run_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

# ----------------------------------------------------------------------
# Additional STATS
# ----------------------------------------------------------------------
ko_assigned = int(df['ko'].notna().sum()) if 'ko' in df.columns else 0

if 'pathways' in df.columns:
    pathway_assigned = int(df['pathways'].apply(lambda x: isinstance(x, list) and len(x) > 0).sum())
else:
    pathway_assigned = 0

category_counts_total = df['cat'].value_counts(dropna=False).to_dict() if 'cat' in df.columns else {}
category_counts_filtered = df_filtered['cat'].value_counts(dropna=False).to_dict() if 'cat' in df_filtered.columns else {}

network_used_fallback = False
shared_pathway_edges = 0
fallback_edges = 0

for _, _, data in G.edges(data=True):
    shared = data.get("shared_pathways", [])
    if isinstance(shared, list) and len(shared) > 0:
        shared_pathway_edges += 1
    else:
        fallback_edges += 1

if G.number_of_edges() > 0 and shared_pathway_edges == 0 and fallback_edges > 0:
    network_used_fallback = True

external_core_matches = int(df['external_core_match'].sum()) if 'external_core_match' in df.columns else 0
external_core_match_types = (
    df['external_core_match_type'].value_counts(dropna=False).to_dict()
    if 'external_core_match_type' in df.columns else {}
)

# ----------------------------------------------------------------------
# MAIN report DICTIONARY
# ----------------------------------------------------------------------
report_dict = {
    "sample_id": sample_id,
    "timestamp": run_timestamp,
    "mode": mode,
    "external_core_ref": external_core_ref,
    "external_core_ref_applied_to_filter": external_core_ref_applied_to_filter,
    "external_core_matches": external_core_matches,
    "external_core_match_types": external_core_match_types,
    "annotation_mode": annotation_mode,
    "library_mode": lib_mode,
    "interpretation": "Expresión (RNA-seq)" if lib_mode == 'R' else "Abundancia normalizada (WGS/metagenoma)",
    "genes_total": int(len(df)),
    "genes_filtered": int(len(df_filtered)),
    "seed_genes": int(seed_gene_count),
    "context_expanded_genes": int(context_expanded_gene_count),
    "genes_with_ko": ko_assigned,
    "genes_with_pathways": pathway_assigned,
    "nodes": int(G.number_of_nodes()),
    "edges": int(G.number_of_edges()),
    "shared_pathway_edges": int(shared_pathway_edges),
    "fallback_edges": int(fallback_edges),
    "network_used_fallback": network_used_fallback,
    "category_counts_total": category_counts_total,
    "category_counts_filtered": category_counts_filtered,
    "external_core_expr_band_counts": df.loc[df['external_core_match'] == True,
    'external_core_expr_band'].value_counts(dropna=False).to_dict() if ('external_core_match' in df.columns and
    'external_core_expr_band' in df.columns) else {},
    "files_generated": {
        "full_table": full_csv,
        "filtered_table": filt_csv,
        "nodes_table": nodes_csv,
        "edges_table": edges_csv,
        "graphml": graphml_out,
        "gexf": gexf_out,
        "gene_ko_pathway": gene_ko_path_csv,
        "network_png": png_out,
        "external_core_matches_table": core_matches_csv
    }
}

# ----------------------------------------------------------------------
# SAVE JSON
# ----------------------------------------------------------------------
json_out = os.path.join(OUT_DIR, f"{sample_id}_SELECTNET_mode_{mode}_run_report.json")
with open(json_out, "w", encoding="utf-8") as f:
    json.dump(report_dict, f, indent=4, ensure_ascii=False)

# ----------------------------------------------------------------------
# SAVE Markdown legible
# ----------------------------------------------------------------------
md_out = os.path.join(OUT_DIR, f"{sample_id}_SELECTNET_mode_{mode}_run_report.md")

md_lines = [
    "# SELECT-NET Run Report",
    "",
    f"- **External core reference:** {external_core_ref}",
    f"- **External core reference applied to filter:** {external_core_ref_applied_to_filter}",
    f"- **External core matches:** {external_core_matches}",
    "",
    f"**Sample ID:** {sample_id}",
    f"**Date/Time:** {run_timestamp}",
    f"**Mode:** {mode}",
    f"**Annotation mode:** {annotation_mode}",
    f"**Library mode:** {lib_mode}",
    f"**Interpretation:** {'Expresión (RNA-seq)' if lib_mode == 'R' else 'Abundancia normalizada (WGS/metagenoma)'}",
    "",
    "## Summary statistics",
    "",
    f"- **Total genes:** {len(df)}",
    f"- **Filtered genes:** {len(df_filtered)}",
    f"- **Genes with KO:** {ko_assigned}",
    f"- **Genes with pathways:** {pathway_assigned}",
    f"- **Nodes:** {G.number_of_nodes()}",
    f"- **Edges:** {G.number_of_edges()}",
    f"- **Shared-pathway edges:** {shared_pathway_edges}",
    f"- **Fallback edges:** {fallback_edges}",
    f"- **Fallback-only network:** {network_used_fallback}",
    "",
    "## External core matching types",
    ""
]

for k, v in external_core_match_types.items():
    md_lines.append(f"- **{k}**: {v}")

md_lines.extend([
    "",
    "## Category distribution (total)",
    ""
])

for k, v in category_counts_total.items():
    md_lines.append(f"- **{k}**: {v}")

md_lines.extend([
    "",
    "## Category distribution (filtered)",
    ""
])

for k, v in category_counts_filtered.items():
    md_lines.append(f"- **{k}**: {v}")

md_lines.extend([
    "",
    "## Generated files",
    "",
    f"- Full table: `{full_csv}`",
    f"- Filtered table: `{filt_csv}`",
    f"- Nodes table: `{nodes_csv}`",
    f"- Edges table: `{edges_csv}`",
    f"- GraphML: `{graphml_out}`",
    f"- GEXF: `{gexf_out}`",
    f"- Gene→KO→Pathway table: `{gene_ko_path_csv}`",
    f"- Network PNG: `{png_out}`",
    f"- External core matches table: `{core_matches_csv}`",
    "",
    "## Publication-ready note",
    "",
    f"This run was generated with SELECT-NET using mode **{mode}** on sample **{sample_id}**. "
    f"For mode C, core-like filtering was driven by the external reference **{external_core_ref}**. "
    f"The workflow exported full and filtered annotation tables, enriched node and edge tables, "
    f"GraphML and GEXF network files, a gene-to-KO-to-pathway mapping table, and a static PNG representation.",
    ""
])

with open(md_out, "w", encoding="utf-8") as f:
    f.write("\n".join(md_lines))

# ----------------------------------------------------------------------
# SAVE main CSV acumulativo
# ----------------------------------------------------------------------
master_csv = os.path.join(OUT_DIR, "SELECTNET_run_summary_master.csv")

master_row = {
    "sample_id": sample_id,
    "timestamp": run_timestamp,
    "mode": mode,
    "external_core_ref": external_core_ref,
    "external_core_ref_applied_to_filter": external_core_ref_applied_to_filter,
    "external_core_matches": external_core_matches,
    "annotation_mode": annotation_mode,
    "library_mode": lib_mode,
    "interpretation": "Expresión (RNA-seq)" if lib_mode == 'R' else "Abundancia normalizada (WGS/metagenoma)",
    "genes_total": int(len(df)),
    "genes_filtered": int(len(df_filtered)),   #compatibilidade historica eles sao os expanded nodes
    "seed_genes": seed_gene_count,
    "context_expanded_genes": context_expanded_gene_count,
    "genes_with_ko": ko_assigned,
    "genes_with_pathways": pathway_assigned,
    "nodes": int(G.number_of_nodes()),
    "edges": int(G.number_of_edges()),
    "external_core_expr_high": int(((df['external_core_match'] == True) & (df['external_core_expr_band'] == 'high')).sum()) if ('external_core_match' in df.columns and 'external_core_expr_band' in df.columns) else 0,
"external_core_expr_low": int(((df['external_core_match'] == True) & (df['external_core_expr_band'] == 'low')).sum()) if ('external_core_match' in df.columns and 'external_core_expr_band' in df.columns) else 0,
    "shared_pathway_edges": int(shared_pathway_edges),
    "fallback_edges": int(fallback_edges),
    "network_used_fallback": network_used_fallback
}

df_master_row = pd.DataFrame([master_row])

if os.path.exists(master_csv):
    df_master_existing = pd.read_csv(master_csv)
    df_master_all = pd.concat([df_master_existing, df_master_row], ignore_index=True)
else:
    df_master_all = df_master_row

df_master_all.to_csv(master_csv, index=False)

print(f"Reporte JSON exportado a: {json_out}")
print(f"Reporte Markdown exportado a: {md_out}")
print(f"Resumen maestro acumulativo exportado a: {master_csv}")

# ==============================================================================
# SOFTWARE VERSIONS AND CITATIONS
# ==============================================================================

#import os
#import sys
#import subprocess

#versions_out = os.path.join(OUT_DIR, "SELECT_NET_versions_and_citations.txt")

#def capture(cmd):
#    try:
#        res = subprocess.run(cmd, capture_output=True, text=False, check=False)
#        raw = res.stdout if res.stdout else res.stderr

#        if raw is None:
#            txt = "[no output]\n"
#        else:
#            try:
#                txt = raw.decode("utf-8")
#            except UnicodeDecodeError:
#                txt = raw.decode("utf-8", errors="replace")

#        header = f"\n[COMMAND] {' '.join(cmd)}\n"
#        return header + txt

#    except Exception as e:
#        return f"\n[COMMAND] {' '.join(cmd)}\n[ERROR] {e}\n"

#report = []
#report.append("SELECT-NET software environment report\n")
#report.append("Generated directly from the active runtime environment.\n")

# Python
#report.append(capture([sys.executable, "--version"]))

# Python packages
#report.append(capture([sys.executable, "-m", "pip", "show", "biopython"]))
#report.append(capture([sys.executable, "-m", "pip", "show", "pandas"]))
#report.append(capture([sys.executable, "-m", "pip", "show", "networkx"]))
#report.append(capture([sys.executable, "-m", "pip", "show", "matplotlib"]))
#report.append(capture([sys.executable, "-m", "pip", "show", "requests"]))

# Main tools
#report.append(capture(["cutadapt", "--version"]))
#report.append(capture(["trim_galore", "--version"]))
#report.append(capture(["fastqc", "--version"]))
#report.append(capture(["samtools", "--version"]))
#report.append(capture(["minimap2", "--version"]))
#report.append(capture(["megahit", "--version"]))
#report.append(capture(["spades", "--version"]))

# Prokka
#report.append(capture([PROKKA_BIN, "--version"]))
#report.append(capture([PROKKA_BIN, "--citation"]))
#report.append(capture(["git", "-C", PROKKA_DIR, "rev-parse", "HEAD"]))
#report.append(capture(["git", "-C", PROKKA_DIR, "describe", "--tags", "--always"]))

# KOfamScan
#report.append(capture([os.path.join(KOFAM_SCAN_DIR, "exec_annotation"), "--version"]))
#report.append(capture(["git", "-C", KOFAM_SCAN_DIR, "rev-parse", "HEAD"]))
#report.append(capture(["git", "-C", KOFAM_SCAN_DIR, "describe", "--tags", "--always"]))

# featureCounts / Subread
#report.append(capture(["dpkg", "-s", "subread"]))

# SRA Toolkit
#report.append(capture(["prefetch", "--version"]))
#report.append(capture(["fasterq-dump", "--version"]))

# Manual citation notes
#report.append("\n[MANUAL CITATION NOTES]\n")
#report.append("Biopython recommended citation:\n")
#report.append(
#    "Cock PJA et al. Biopython: freely available Python tools for computational "
#    "molecular biology and bioinformatics. Bioinformatics. 2009;25(11):1422-1423. "
#    "https://doi.org/10.1093/bioinformatics/btp163\n\n"
#)

#report.append("KOfamScan / KofamKOALA associated citation:\n")
#report.append(
#    "Aramaki T et al. KofamKOALA: KEGG Ortholog assignment based on profile HMM "
#    "and adaptive score threshold. Bioinformatics. 2020;36(7):2251-2252. "
#    "https://doi.org/10.1093/bioinformatics/btz859\n\n"
#)

#with open(versions_out, "w", encoding="utf-8") as f:
#    f.write("".join(report))

#print(f"Version and citation report exported to: {versions_out}")
#print("Exists?", os.path.exists(versions_out))
#print("File size (bytes):", os.path.getsize(versions_out) if os.path.exists(versions_out) else "NOT FOUND")