#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright © 2019 yech <yech1990@gmail.com>
#
# Distributed under terms of the MIT license.

"""Chromatogram File Utils.

- update in 20190405
"""

import re
import sys
from datetime import date


class CountMutations:
    def __init__(self):
        pass

    def _to_int(self, str):
        """Convert str to int if it can otherwise 0."""
        tmp = 1
        try:
            tmp = int(str)
        except ValueError:
            pass

        return tmp

    def _get_date(self, seq_def):
        """Parses a sequence definition and extracts the date."""
        parts = seq_def.split("|")

        # Probably no date if there is no | in the line
        if len(parts) == 1:
            return date(1800, 1, 1)

        try:
            y = int(parts[-4])
        except ValueError:
            print("Failed on this definition line")
            sys.exit(-1)
            m = self._to_int(parts[-3])
            d = self._to_int(parts[-2])
            dt = date(y, m, d)
            return dt

    def _get_gi(self, seq_def):
        """Parses genbank id out."""
        parts = seq_def.split("|")
        if len(parts) > 1:
            return parts[1]
        else:
            return seq_def

    def _get_name(self, seq_def):
        """Parses a sequence definition and extracts the name Assumed to be in
        the 2 column after split( '|' ) and between ()"""
        p = re.compile(r"\((.*)\)")
        parts = seq_def.split("|")
        if len(parts) > 1:
            m = p.search(parts[2])
            return m.group(1)
        else:
            return seq_def

    def parse(self, mutations_file, cutoff_date):
        """Parse a mutations file generated by mutalign.py."""
        fh = open(mutations_file)
        chart = []
        mut_count = 0
        last_date = None
        inc = 1
        for line in fh:
            # Start new date and set counter to 0
            if not line.startswith("Q: "):
                # Set mutation count for last date
                if last_date:
                    chart[-1][3] = mut_count

                # Get the new date and set it as last_date
                last_date = self._get_date(line)

                if last_date < cutoff_date:
                    inc = -1
                else:
                    inc = 1
                name = self._get_name(line)
                gi = self._get_gi(line)
                chart.append([gi, last_date, name, 0])
                mut_count = 0
            # Else Count mutations
            mut_count += inc

        fh.close()

        chart[-1][3] = mut_count

        return chart

    def get_chart(self, mutations_file, cutoff_date):
        p = self.parse(mutations_file, cutoff_date)
        for gi, dt, name, num in p:
            print("%s,%s,%s,%s" % (gi, name, dt, num))


def parse_date(dte):
    if "-" in dte:
        p = dte.split("-")
    elif "/" in dte:
        p = dte.split("/")

    if len(p) == 3:
        return date(int(p[0]), int(p[1]), int(p[2]))
