import csv
import itertools
import logging
import pysam
import random
import tabix
import sys

from collections import defaultdict

from ..utils import NGSFile, NGSFileType, GFF

logger = logging.getLogger('strandedness')


def get_filtered_reads_from_region(samfile,
                                   gene,
                                   min_quality=30,
                                   apply_filters=True):
    for read in samfile.fetch(gene['seqname'], gene["start"], gene["end"]):
        if apply_filters and (read.is_qcfail or read.is_duplicate
                              or read.is_secondary or read.is_unmapped
                              or read.mapq < min_quality):
            continue
        yield read


def disqualify_gene(gene, gff_tabix, only_consider_protein_genes=True):
    # potentially only consider protein coding genes.
    if only_consider_protein_genes:
        if not "attr_gene_type" in gene:
            return True
        if not "protein" in gene['attr_gene_type']:
            return True

    # if there are overlapping features on the positive and negative strand
    # ignore this gene.
    hits = gff_tabix.query(gene['seqname'], gene['start'], gene['end'])
    has_positive_gene = None
    has_negative_gene = None

    for hit in hits:
        # must be a gene
        if not hit[2] == "gene":
            continue

        if hit[6] == "+":
            has_positive_gene = hit
        elif hit[6] == "-":
            has_negative_gene = hit

        if has_positive_gene and has_negative_gene:
            break

    if has_positive_gene and has_negative_gene:
        return True

    return False

def get_reads_rg(read, default="unknown"):
  for (k, v) in read.tags:
    if k == "RG":
      return v
    
  return default

def get_predicted_strandedness(forward_evidence_pct, reverse_evidence_pct):
    predicted = "Inconclusive"
    if 0.4 <= forward_evidence_pct and forward_evidence_pct <= 0.6:
        predicted = "Unstranded"
    if 0.8 <= forward_evidence_pct:
        predicted = "Stranded-Forward"
    elif 0.8 <= reverse_evidence_pct:
        predicted = "Stranded-Reverse"
    
    return predicted

def main(ngsfiles,
         gene_model_file,
         outfile=sys.stdout,
         delimiter="\t",
         n_genes=100,
         minimum_reads_per_gene=10,
         only_protein_coding_genes=True,
         min_mapq=30,
         split_by_rg=False):
    logger.info("Arguments:")
    logger.info("  - Gene model file: {}".format(gene_model_file))
    logger.info("  - Number of genes: {}".format(n_genes))
    logger.info("  - Minimum reads per gene: {}".format(minimum_reads_per_gene))
    logger.info("  - Only consider protein coding genes: {}".format(only_protein_coding_genes))
    logger.info("  - Minimum MAPQ: {}".format(min_mapq))
    logger.info("  - Split by RG: {}".format(split_by_rg))

    logger.info("Reading gene model...")
    gff = GFF(gene_model_file, feature_type="gene")
    logger.info("  - {} features processed.".format(len(gff.entries)))
    gff_tabix = tabix.open(gene_model_file)
    logger.info("  - Tabix loaded for feature lookup.")

    fieldnames = [ "TotalReads", "ForwardPct", "ReversePct", "Predicted"] 
    if split_by_rg:
      fieldnames = [ "ReadGroup" ] + fieldnames 
    fieldnames = [ "File" ] + fieldnames

    writer = csv.DictWriter(outfile,
                            fieldnames=fieldnames,
                            delimiter=delimiter)
    writer.writeheader()
    outfile.flush()

    for ngsfilepath in ngsfiles:
        logger.info("Processing {}...".format(ngsfilepath))
        try:
            ngsfile = NGSFile(ngsfilepath)
        except:
            result = {
                "File": ngsfilepath,
                "TotalReads": "N/A",
                "ForwardPct": "N/A",
                "ReversePct": "N/A",
                "Predicted": "File not found." 
            }

            if args.split_by_rg:
              result["ReadGroup"] = "N/A"

            writer.writerow(result)
            outfile.flush()
            continue

        if ngsfile.filetype != NGSFileType.BAM:
            raise RuntimeError("Invalid file: {}. `strandedness` only supports aligned BAM files!".format(ngsfilepath))
        samfile = ngsfile.handle

        n_tested_genes = 0
        n_reads_observed = 0
        gene_blacklist = set()
        read_groups = [rg['ID'] for rg in samfile.header['RG']] + ["unknown"]
        overall_evidence = defaultdict(dict)
        for rg in read_groups:
            overall_evidence[rg] = defaultdict(int)

        logger.debug("Starting sampling...")
        while True:
            if n_tested_genes >= n_genes:
                break

            gene = random.choice(gff.entries)

            if gene["attr_gene_id"] in gene_blacklist:
                continue

            if disqualify_gene(gene, gff_tabix, only_consider_protein_genes=only_protein_coding_genes):
                continue

            logging.debug("== Candidate Gene ==")
            logging.debug("  [*] Name: {}".format(gene['attr_gene_name']))
            logging.debug("  [*] Location: {}:{}-{} ({})".format(
                gene['seqname'], gene['start'], gene['end'], gene['strand']))
            logging.debug("  [*] Actions:")

            gene_blacklist.add(gene["attr_gene_id"])
            relevant_reads = get_filtered_reads_from_region(samfile, gene, min_quality=min_mapq)

            reads_in_gene = 0
            this_genes_evidence = defaultdict(dict)
            for rg in read_groups:
                this_genes_evidence[rg] = defaultdict(int)

            for read in relevant_reads:
                reads_in_gene += 1
                if not read.is_paired:
                    raise RuntimeError(
                        "This tool currently only works for paired-end data! Please contact the author if you'd like SE data to be supported."
                    )

                if read.is_read1:
                    read_id = "1"
                elif read.is_read2:
                    read_id = "2"
                else:
                    raise RuntimeError("Read is not read 1 or read 2?")

                if not read.is_reverse:
                    read_strand_id = "+"
                else:
                    read_strand_id = "-"

                gene_strand_id = gene['strand']
                reads_observed_state = read_id + read_strand_id + gene_strand_id
                this_reads_rg = get_reads_rg(read)
                this_genes_evidence[this_reads_rg][reads_observed_state] += 1

            if reads_in_gene >= minimum_reads_per_gene:
                logging.debug("    - Sufficient read count ({} >= {})".format(
                    reads_in_gene, minimum_reads_per_gene))
                logging.debug("    - {}".format(" ".join([
                    "{}:{}".format(k, v)
                    for k, v in this_genes_evidence.items()
                ])))

                for rg in this_genes_evidence.keys():
                    for state in this_genes_evidence[rg].keys():
                        overall_evidence[rg][state] += this_genes_evidence[rg][state]

                n_tested_genes += 1
                n_reads_observed += reads_in_gene
            else:
                logging.debug("    - Read count too low ({} < {})".format(
                    reads_in_gene, minimum_reads_per_gene))

        if split_by_rg:
          for rg in read_groups:
            evidence_stranded_forward = overall_evidence[rg]["1++"] + overall_evidence[rg][ "1--"] + overall_evidence[rg]["2+-"] + overall_evidence[rg]["2-+"]
            evidence_stranded_reverse = overall_evidence[rg]["1+-"] + overall_evidence[rg][ "1-+"] + overall_evidence[rg]["2++"] + overall_evidence[rg]["2--"]
            total_reads = evidence_stranded_forward + evidence_stranded_reverse
            if total_reads == 0 and rg == "unknown":
              continue

            forward_pct = -1 if total_reads == 0 else round(evidence_stranded_forward / total_reads, 4)
            reverse_pct = -1 if total_reads == 0 else round(evidence_stranded_reverse / total_reads, 4)
            predicted = get_predicted_strandedness(forward_pct, reverse_pct)

            result = {
                "File": ngsfilepath,
                "ReadGroup": rg,
                "TotalReads": total_reads,
                "ForwardPct": forward_pct,
                "ReversePct": reverse_pct,
                "Predicted": predicted
            }

            writer.writerow(result)
            outfile.flush()
        else:
          evidence_stranded_forward = 0
          evidence_stranded_reverse = 0

          for rg in read_groups:
            evidence_stranded_forward += overall_evidence[rg]["1++"] + overall_evidence[rg][ "1--"] + overall_evidence[rg]["2+-"] + overall_evidence[rg]["2-+"]
            evidence_stranded_reverse += overall_evidence[rg]["1+-"] + overall_evidence[rg][ "1-+"] + overall_evidence[rg]["2++"] + overall_evidence[rg]["2--"]

          total_reads = evidence_stranded_forward + evidence_stranded_reverse
          forward_pct = round(evidence_stranded_forward / total_reads, 4)
          reverse_pct = round(evidence_stranded_reverse / total_reads, 4)
          predicted = get_predicted_strandedness(forward_pct, reverse_pct)

          result = {
              "File": ngsfilepath,
              "TotalReads": total_reads,
              "ForwardPct": forward_pct,
              "ReversePct": reverse_pct,
              "Predicted": predicted
          }

          writer.writerow(result)
          outfile.flush()
