"""Ribodesigner module"""

import sys
import pandas as pd
import pysam
import subprocess
from sequana import logger
from pathlib import Path

logger.setLevel("INFO")


def reverse_complement(sequence):
    """Reverse complement a nucleotide sequence."""

    trans = sequence.maketrans("ACGT", "TGCA")
    rev_comp_sequence = sequence.translate(trans)[::-1]

    return rev_comp_sequence


def get_rna_pos_from_gff(gff, seq_type="rRNA"):
    """Convert a GFF file into a pandas DataFrame filtered according to the
    seq_type.

    :param gff: GFF annotation file
    :param seq_type: string describing sequence annotation type (column 3 in GFF) to
    select rRNA from.
    """

    df = pd.read_csv(
        gff,
        sep="\t",
        comment="#",
        names=["seqid", "source", "type", "start", "end", "score", "strand", "phase", "attributes"],
    )
    df = df.query("type == @seq_type")

    if df.empty:
        raise IOError(f"No '{seq_type}' type was found in {gff}.")

    logger.info(f"Found {df.shape[0]} '{seq_type}' entries in annotation file.")

    return df


def extract_regions_from_fasta(fasta_in, regions_df, fasta_out):
    """From a genomic FASTA file fasta_in, extract regions stored in regions_df
    dataframe obtained with get_rna_pos_from_gff. Export the corresponding
    regions to a FASTA file fasta_out.

    :param fasta_in: path to genomic FASTA file to extract rRNA regions from.
    :param regions_df: pandas DataFrame generated by get_rna_pos_from_gff
    :param fasta_out: path to FASTA file where rRNA regions will be exported to.
    """

    count = 0

    with pysam.Fastafile(fasta_in) as fas:
        with open(fasta_out, "w") as fas_out:
            for rna in regions_df.to_dict("records"):
                region = f"{rna['seqid']}:{rna['start']}-{rna['end']}"
                seq_record = f">{rna['seqid']}:{rna['start']}-{rna['end']}\n{fas.fetch(region=region)}\n"
                fas_out.write(seq_record)
                count += 1

    logger.info(f"{count} regions were extracted from '{fasta_in}' to '{fasta_out}'")


def get_probes(fasta_in, fasta_out, probe_len=50, inter_probe_space=15):
    """From a FASTA file of nucleotide sequence to be targetted, generate a FASTA
    file with probe designed of size 'probe_len' and with a space between probes
    of 'inter_probe_space'

    :param fasta_in: path to FASTA file with extracted rRNA regions.
    :param fasta_out: path to FASTA file where probes will be exported to.
    :param prob_len: the size for probes in nucleotides.
    :param inter_probe_space: the space between probes in nucleotides.
    """

    count = 0
    with open(fasta_out, "w") as fas_out:

        with pysam.FastxFile(fasta_in) as fas:
            for seq in fas:

                for start in range(0, len(seq.sequence) - probe_len, probe_len + inter_probe_space):
                    count += 2

                    end = start + probe_len

                    probe_fw = f">{seq.name}:probe_{start}_{end}_forward\n{seq.sequence[start:end]}\n"
                    rev_seq = reverse_complement(seq.sequence)
                    probe_rv = f">{seq.name}:probe_{start}_{end}_reverse\n{rev_seq[start:end]}\n"

                    fas_out.write(probe_fw + probe_rv)

    logger.info(f"{count} probes designed.")


def cluster_probes(fasta_in, fasta_out, seq_id_thres=0.80, threads=4):
    """Use cd-hit-est to cluster highly similar probes

    :param fasta_in: path to FASTA file with probes.
    :param fasta_out: path to FASTA file with probes after filtering.
    """
    log_file = Path(fasta_out).with_suffix(".log")

    cmd = f"cd-hit-est -i {fasta_in} -o {fasta_out} -c {seq_id_thres} -n {threads}"
    logger.info(f"Clustering probes with command: {cmd} (log in '{log_file}').")

    with open(log_file, "w") as f:
        subprocess.run(cmd, shell=True, check=True, stdout=f)


def fasta_to_csv(fasta_in, csv_out):
    """Convert a FASTA file into a CSV file with the oligo name in column 1 and the
    oligo sequence in column 2

    :param fasta_in: path to FASTA file with clustered probes.
    :param fasta_out: path to CSV file with clustered probes.
    """

    seq_dict = {"id": [], "sequence": []}

    with pysam.FastxFile(fasta_in) as fas:
        for n, seq in enumerate(fas):
            seq_dict["id"].append(seq.name)
            seq_dict["sequence"].append(seq.sequence)

    logger.info(f"After clustering, keeping {n} probes.")
    logger.info(f"Creating probes csv table '{csv_out}'.")
    pd.DataFrame(seq_dict).to_csv(csv_out, index=False)

Add ribodesigner library