#!/usr/bin/env python
# -*- coding: UTF-8 -*-

"""
Use genetic map to break chimeric scaffolds, order and orient scaffolds onto
chromosomes.
"""
import os.path as op
import sys
import logging
import numpy as np

from collections import Counter
from functools import lru_cache
from itertools import combinations, groupby

from jcvi.formats.base import BaseFile, LineFile, must_open, read_block
from jcvi.formats.bed import Bed, fastaFromBed
from jcvi.apps.base import OptionParser, ActionDispatcher, need_update


MSTheader = """population_type {0}
population_name LG
distance_function kosambi
cut_off_p_value 0.000001
no_map_dist 10.0
no_map_size 0
missing_threshold {1}
estimation_before_clustering no
detect_bad_data yes
objective_function ML
number_of_loci {2}
number_of_individual {3}
"""


class BinMap(BaseFile, dict):
    def __init__(self, filename):
        super(BinMap, self).__init__(filename)

        fp = open(filename)
        for header, seq in read_block(fp, "group "):
            lg = header.split()[-1]
            self[lg] = []
            for s in seq:
                if s.strip() == "" or s[0] == ";":
                    continue
                marker, pos = s.split()
                pos = int(float(pos) * 1000)
                self[lg].append((marker, pos))

    def print_to_bed(self, filename="stdout", switch=False, sep="."):
        """Print the genetic map in the BED format.

        Args:
            filename (str, optional): Output filename. Defaults to "stdout".
            switch (bool, optional): Use linkage group as seqid. Defaults to False.
            sep (str, optional): Separator that delimits scaffold name and position. Defaults to ".".
        """
        fw = must_open(filename, "w")
        for lg, markers in sorted(self.items()):
            for marker, pos in markers:
                if not switch:
                    line = (lg, pos, pos + 1, marker)
                else:
                    seqid_spos = marker.rsplit(sep, 1)
                    if len(seqid_spos) != 2:
                        logging.error(
                            "Error: `{}` must be in the form e.g. `name{}position`".format(
                                marker, sep
                            )
                        )
                        continue
                    seqid, spos = seqid_spos
                    spos = int(spos)
                    marker = "{0}:{1}".format(lg, pos / 1000.0)
                    line = (seqid, spos - 1, spos, marker)
                print("\t".join(str(x) for x in line), file=fw)
        fw.close()


class MSTMapLine(object):
    def __init__(self, row, startidx=3):
        args = row.split()
        self.id = args[0]
        self.seqid, pos = self.id.split(".")
        self.pos = int(pos)
        self.genotype = "".join(args[startidx:])

    def __len__(self):
        return len(self.genotype)

    def __str__(self):
        return "{0}: {1}".format(self.id, self.genotype)

    @property
    def bedline(self):
        return "\t".join(str(x) for x in (self.seqid, self.pos - 1, self.pos, self.id))


class MSTMap(LineFile):
    def __init__(self, filename):
        super(MSTMap, self).__init__(filename)
        fp = open(filename)
        startidx = 1
        for row in fp:
            if row.startswith("locus_name"):
                if row.split()[1] == "seqid":
                    startidx = 3
                self.header = row.split()
                break

        for row in fp:
            self.append(MSTMapLine(row, startidx=startidx))

        self.nmarkers = len(self)
        self.nind = len(self[0].genotype)
        logging.debug(
            "Map contains {0} markers in {1} individuals".format(
                self.nmar
kers, self.nind
            )
        )


class MSTMatrix(object):
    def __init__(self, matrix, markerheader, population_type,
 missing_threshold):
        self.matrix = matrix
        self.markerheader = markerheader
        self.population_type = population_type
        self.missing_threshold = missing_threshold
        self.ngenotypes = len(matrix)
        self.nind = len(markerheader) - 1
        assert self.nind == len(matrix[0]) - 1
        logging.debug(
            "Imported {0} markers and {1} individuals.".format(
                self.ngenotypes, self.nind
            )
        )

    def write(self, filename="stdout", header=True):
        fw = must_open(filename, "w")
        if header:
            print(
                MSTheader.format(
                    self.population_type,
                    self.missing_threshold,
                    self.ngenotypes,
                    self.nind,
                ),
                file=fw,
            )
        print("\t".join(self.markerheader), file=fw)
        for m in self.matrix:
            print("\t".join(m), file=fw)


def main():

    actions = (
        ("breakpoint", "find scaffold breakpoints using genetic map"),
        ("ld", "calculate pairwise linkage disequilibrium"),
        ("bed", "convert MSTmap output to bed format"),
        ("fasta", "extract markers based on map"),
        ("anchor", "anchor scaffolds based on map"),
        ("rename", "rename markers according to the new mapping locations"),
        ("header", "rename lines in the map header"),
        # Plot genetic map
        ("blat", "make ALLMAPS input csv based on sequences"),
        ("dotplot", "make dotplot between chromosomes and linkage maps"),
    )
    p = ActionDispatcher(actions)
    p.dispatch(globals())


def blat(args):
    """
    %prog blat map1.txt ref.fasta

    Make ALLMAPS input csv based on sequences. The tab-delimited txt file
    include: name, LG, position, sequence.
    """
    from jcvi.formats.base import is_number
    from jcvi.formats.blast import best as blast_best, bed as blast_bed
    from jcvi.apps.align import blat as blat_align

    p = OptionParser(blat.__doc__)
    _, args = p.parse_args(args)

    if len(args) != 2:
        sys.exit(not p.print_help())

    maptxt, ref = args
    pf = maptxt.rsplit(".", 1)[0]
    register = {}
    fastafile = pf + ".fasta"
    fp = open(maptxt)
    fw = open(fastafile, "w")
    for row in fp:
        name, lg, pos, seq = row.split()
        if not is_number(pos):
            continue
        register[name] = (pf + "-" + lg, pos)
        print(">{0}\n{1}\n".format(name, seq), file=fw)
    fw.close()

    blatfile = blat_align([ref, fastafile])
    bestfile = blast_best([blatfile])
    bedfile = blast_bed([bestfile])
    b = Bed(bedfile).order

    pf = ".".join((op.basename(maptxt).split(".")[0], op.basename(ref).split(".")[0]))
    csvfile = pf + ".csv"
    fp = open(maptxt)
    fw = open(csvfile, "w")
    for row in fp:
        name, lg, pos, seq = row.split()
        if name not in b:
            continue
        bbi, bb = b[name]
        scaffold, scaffold_pos = bb.seqid, bb.start
        print(",".join(str(x) for x in (scaffold, scaffold_pos, lg, pos)), file=fw)
    fw.close()


def dotplot(args):
    """
    %prog dotplot map.csv ref.fasta

    Make dotplot between chromosomes and linkage maps.
    The input map is csv formatted, for example:

    ScaffoldID,ScaffoldPosition,LinkageGroup,GeneticPosition
    scaffold_2707,11508,1,0
    scaffold_2707,11525,1,1.2
    """
    from natsort import natsorted
    from jcvi.assembly.allmaps import CSVMapLine
    from jcvi.formats.sizes import Sizes
    from jcvi.graphics.base import shorten
    from jcvi.graphics.dotplot import (
        plt,
        savefig,
        markup,
        normalize_axes,
        downsample,
        plot_breaks_and_labels,
        thousands,
    )

    p = OptionParser(dotplot.__doc__)
    p.set_outfile(outfile=None)
    opts, args, iopts = p.set_image_options(
        args, figsize="8x8", style="dark", dpi=90, cmap="copper"
    )

    if len(args) != 2:
        sys.exit(not p.print_help())

    csvfile, fastafile = args
    sizes = natsorted(Sizes(fastafile).mapping.items())
    seen = set()
    raw_data = []

    fig = plt.figure(1, (iopts.w, iopts.h))
    root = fig.add_axes([0, 0, 1, 1])  # the whole canvas
    ax = fig.add_axes([0.1, 0.1, 0.8, 0.8])  # the dot plot

    fp = m