"""
    This file is part of Picopore.

    Picopore is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    Picopore is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Picopore.  If not, see <http://www.gnu.org/licenses/>.
"""
import subprocess
import numpy as np
import h5py
import os
from numpy.lib.recfunctions import drop_fields, append_fields
from functools import partial

from picopore.util import log, isGroup, getDtype, findDatasets, rewriteDataset, recursiveCollapseGroups, uncollapseGroups, getPrefixedFilename

__basegroup_name__ = "Picopore"
__raw_compress_keywords__ = ["Alignment","Log","Configuration","HairpinAlign","Calibration_Strand","Hairpin_Split","EventDetection","Events","Segmentation"]
__raw_compress_summary__ = ["Summary"]
__raw_compress_fastq__ = ["BaseCalled"]
__raw_compress_fastq_summary__ = ["Basecall"]

def chooseCompressFunc(revert, mode, fastq, summary, manual, realtime=False):
    name = "Performing "
    if realtime:
        name += "real time "
    if revert:
        if mode == 'lossless':
            func = losslessDecompress
            name += "lossless decompression"
        elif mode == 'deep-lossless':
            func = deepLosslessDecompress
            name += "deep lossless decompression"
        else:
            log("Unable to revert raw files. Please use a basecaller instead.")
            exit(1)
    else:
        if mode == 'lossless':
            func = losslessCompress
            name += "lossless compression"
        elif mode == 'deep-lossless':
            func = deepLosslessCompress
            name += "deep lossless compression"
        elif mode == 'raw':
            name += "raw compression "
            if manual is not None:
                name += "with manual keyword " + manual
                keywords = [manual]
            else:
                keywords = __raw_compress_keywords__
                if fastq and summary:
                    name += "with FASTQ and summary"
                elif fastq:
                    keywords += __raw_compress_summary__
                    name += "with FASTQ and no summary"
                elif summary:
                    keywords += __raw_compress_fastq__
                    name += "with summary and no FASTQ"
                else:
                    keywords += __raw_compress_fastq_summary__
                    name += "with no summary and no FASTQ"
            func = partial(rawCompress, keywords=keywords)
    try:
        return partial(compress, func), name
    except NameError:
        log("No compression method selected")
        exit(1)

def indexToZero(f, path, col, name="picopore.{}_index", dataColumn=None):
    dataset = f[path]
    name = name.format(col)
    data = f[path].value
    if not name in dataset.attrs.keys():
        dataColumn = data[col] if dataColumn is None else dataColumn
        start_index = min(dataColumn)
        dataset.attrs.create(name, start_index, dtype=getDtype(start_index))
        dataColumn = dataColumn-start_index
        data = drop_fields(data, [col])
        data = append_fields(data, [col], [dataColumn], [getDtype(dataColumn)])
    return data

def deepLosslessCompress(f, group):
    paths = findDatasets(f, group, "Events")
    paths = [path for path in paths if "Basecall" in path]
    # index event detection
    if "UniqueGlobalKey/channel_id" in f:
        sampleRate = f["UniqueGlobalKey/channel_id"].attrs["sampling_rate"]
        for path in paths:
            if f[path].parent.parent.attrs.__contains__("event_detection"):
                # index back to event detection
                dataset = f[path].value
                start = np.array([int(round(sampleRate * i)) for i in dataset["start"]])
                dataset = indexToZero(f, path, "start", dataColumn=start)
                move = dataset["move"] # rewrite move dataset because it's int64 for max 2
                # otherwise, event by event
                dataset = drop_fields(dataset, ["mean", "stdv", "length", "move"])
                dataset = append_fields(dataset, ["move"], [move], [getDtype(move)])
                rewriteDataset(f, path, compression="gzip", compression_opts=9, dataset=dataset)
                # rewrite eventdetection too - start is also way too big here
                eventDetectionPath = findDatasets(f, "all", entry_point=f[path].parent.parent.attrs.get("event_detection"))[0]
                if "picopore.start_index" not in f[eventDetectionPath].attrs.keys():
                    eventData = indexToZero(f, eventDetectionPath, "start")
                    rewriteDataset(f, eventDetectionPath, compression="gzip", compression_opts=9, dataset=eventData)

    if __basegroup_name__ not in f:
        f.create_group(__basegroup_name__)
        for name, group in f.items():
            if name != __basegroup_name__:
                recursiveCollapseGroups(f, __basegroup_name__, name, group)
    return losslessCompress(f, group)

def deepLosslessDecompress(f, group):
    # rebuild group hierarchy
    if __basegroup_name__ in f.keys():
        uncollapseGroups(f, f[__basegroup_name__])
    paths = findDatasets(f, group)
    paths = [path for path in paths if "Basecall" in path]
    sampleRate = f["UniqueGlobalKey/channel_id"].attrs["sampling_rate"]
    for path in paths:
        if f[path].parent.parent.attrs.__contains__("event_detection"):
            # index back to event detection
            dataset = f[path].value
            if "mean" not in dataset.dtype.names:
                eventDetectionPath = findDatasets(f, "all", entry_point=f[path].parent.parent.attrs.get("event_detection"))[0]
                eventData = f[eventDetectionPath].value
                try:
                    start = eventData["start"] + f[eventDetectionPath].attrs["picopore.start_index"]
                    del f[eventDetectionPath].attrs["picopore.start_index"]
                    eventData = drop_fields(eventData, ["start"])
                    eventData = append_fields(eventData, ["start"], [start], [getDtype(start)])
                    rewriteDataset(f, eventDetectionPath, compression="gzip", compression_opts=1, dataset=eventData)
                except KeyError:
                    # must have been compressed without start indexing
                    pass
                try:
                    start_index = f[path].attrs["picopore.start_index"]
                    del f[path].attrs["picopore.start_index"]
                except KeyError:
                    # must have been compressed without start indexing
                    start_index=0
                start = dataset["start"][0] + start_index
                end = dataset["start"][-1] + start_index
                # constrain to range in basecall
                eventData = eventData[np.logical_and(eventData["start"] >= start, eventData["start"] <= end)]
                # remove missing events
                i=0
                keepIndex = []
                for time in dataset["start"]:
                    while eventData["start"][i] != time + start_index and i < eventData.shape[0]:
                        i += 1
                    keepIndex.append(i)
                eventData = eventData[keepIndex]
                dataset = drop_fields(dataset, "start")
                start = [i/sampleRate for i in eventData["start"]]
                length = [i/sampleRate for i in eventData["length"]]
                dataset = append_fields(dataset, ["mean", "start", "stdv", "length"], [eventData["mean"], start, eventData["stdv"], length])
                rewriteDataset(f, path, dataset=dataset)
    return losslessDecompress(f, group)

def losslessCompress(f, group):
    paths = findDatasets(f, group, keyword="Events")
    paths.extend(findDatasets(f, group, keyword="Alignment"))
    paths.extend(findDatasets(f, "all", keyword="Signal", entry_point="Raw"))
    for path in paths:
        rewriteDataset(f, path, "gzip", 9)
    return "GZIP=9"

def losslessDecompress(f, group):
    paths = findDatasets(f, group, keyword="Events")
    paths.extend(findDatasets(f, group, keyword="Alignment"))
    paths.extend(findDatasets(f, "all", keyword="Signal", entry_point="Raw"))
    for path in paths:
        rewriteDataset(f, path)
    return "GZIP=1"

def rawCompress(f, group, keywords):
    if "Picopore" in f:
        log("{} is compressed using picopore deep-lossless compression. Please use picpore --revert --mode deep-lossless before attempting raw compression.".format(f.filename))
    else:
        paths = []
        for kw in keywords:
            paths.extend(findDatasets(f, group, keyword=kw))
        for path in paths:
            if path in f:
                del f[path]
        try:
            if len(f["Analyses"].keys()) == 0:
                del f["Analyses"]
        except KeyError:
            # no analyses, no worries
            pass
    return "GZIP=9"

def compress(func, filename, group="all"):
    try:
        with h5py.File(filename, 'r+') as f:
            filtr = func(f, group)
        subprocess.call(["h5repack","-f",filtr,filename, "{}.tmp".format(filename)])
        subprocess.call(["mv","{}.tmp".format(filename),filename])
        return os.path.getsize(filename)
    except Exception as e:
        log("ERROR: {} on file {}".format(str(e), filename))
        if os.path.isfile("{}.tmp".format(filename)):
            os.remove("{}.tmp".format(filename))
        return os.path.getsize(filename)
