# Filename: ahrs.py
# -*- coding: utf-8 -*-
# pylint: disable=locally-disabled
"""
AHRS calibration.

"""

import io
from collections import defaultdict
import time
import xml.etree.ElementTree as ET

import km3db

import numpy as np
from numpy import cos, sin, arctan2
import km3pipe as kp
from km3pipe.tools import timed_cache
from km3pipe.io.daq import TMCHData

__author__ = "Tamas Gal"
__email__ = "tgal@km3net.de"
__status__ = "Development"

log = kp.logger.get_logger(__name__)  # pylint: disable=C0103

# log.setLevel("DEBUG")


class AHRSCalibrator(kp.Module):
    """Calculates AHRS yaw, pitch and roll from median A and H of an interval.

    Parameters
    ----------
    det_id: int
        The detector ID, e.g. 29)

    Other Parameters
    ----------------
    interval: int (accumulation interval in [sec], default: 10s)

    Notes
    -----
    Writes 'AHRSCalibration' in the blob with:
        dict: key=dom_id, value=tuple: (timestamp, du, floor, yaw, pitch, roll)

    """

    def configure(self):
        det_id = self.require("det_id")
        self.interval = self.get("interval") or 10  # in sec
        self.A = defaultdict(list)
        self.H = defaultdict(list)
        self.detector = kp.hardware.Detector(det_id=det_id)
        self.clbmap = km3db.CLBMap(det_id)
        self.timestamp = time.time()

    def process(self, blob):
        tmch_data = TMCHData(io.BytesIO(blob["CHData"]))
        dom_id = tmch_data.dom_id
        try:
            du, floor, _ = self.detector.doms[dom_id]
        except KeyError:  # base CLB
            return blob

        self.A[dom_id].append(tmch_data.A)
        self.H[dom_id].append(tmch_data.H)

        if time.time() - self.timestamp > self.interval:
            self.timestamp = time.time()
            calib = self.calibrate()
            blob["AHRSCalibration"] = calib

        return blob

    def calibrate(self):
        """Calculate yaw, pitch and roll from the median of A and H.

        After successful calibration, the `self.A` and `self.H` are reset.
        DOMs with missing AHRS pre-calibration data are skipped.

        Returns
        -------
        dict: key=dom_id, value=tuple: (timestamp, du, floor, yaw, pitch, roll)

        """
        now = time.time()
        dom_ids = self.A.keys()
        print("Calibrating AHRS from median A and H for {} DOMs.".format(len(dom_ids)))
        calibrations = {}
        for dom_id in dom_ids:
            print("Calibrating DOM ID {}".format(dom_id))
            clb_upi = self.clbmap.doms_ids[dom_id].clb_upi
            ahrs_calib = get_latest_ahrs_calibration(clb_upi)
            if ahrs_calib is None:
                log.warning("AHRS calibration missing for '{}'".format(dom_id))
                continue
            du, floor, _ = self.detector.doms[dom_id]
            A = np.median(self.A[dom_id], axis=0)
            H = np.median(self.H[dom_id], axis=0)

            cyaw, cpitch, croll = fit_ahrs(A, H, *ahrs_calib)
            calibrations[dom_id] = (now, du, floor, cyaw, cpitch, croll)
        self.A = defaultdict(list)
        self.H = defaultdict(list)
        return calibrations


def fit_ahrs(A, H, Aoff, Arot, Hoff, Hrot):
    """Calculate yaw, pitch and roll for given A/H and calibration set.

    Author: Vladimir Kulikovskiy

    Parameters
    ----------
    A: list, tuple or numpy.array of shape (3,)
    H: list, tuple or numpy.array of shape (3,)
    Aoff: numpy.array of shape(3,)
    Arot: numpy.array of shape(3, 3)
    Hoff: numpy.array of shape(3,)
    Hrot: numpy.array of shape(3, 3)

    Returns
    -------
    yaw, pitch, roll

    """
    Acal = np.dot(A - Aoff, Arot)
    Hcal = np.dot(H - Hoff, Hrot)

    # invert axis for DOM upside down
    for i in (1, 2):
        Acal[i] = -Acal[i]
        Hcal[i] = -Hcal[i]

    roll = arctan2(-Acal[1], -Acal[2])
    pitch = arctan2(Acal[0], np.sqrt(Acal[1] * Acal[1] + Acal[2] * Acal[2]))
    yaw = arctan2(
        Hcal[2] * sin(roll) - Hcal[1] * cos(roll),
        sum(
            (
                Hcal[0] * cos(pitch),
                Hcal[1] * sin(pitch) * sin(roll),
                Hcal[2] * sin(pitch) * cos(roll),
            )
        ),
    )

    yaw = np.degrees(yaw)
    while yaw < 0:
        yaw += 360
    # yaw = (yaw + magnetic_declination + 360 ) % 360
    roll = np.degrees(roll)
    pitch = np.degrees(pitch)
    return yaw, pitch, roll


@timed_cache(hours=1, maxsize=None, typed=False)
def get_latest_ahrs_calibration(clb_upi, max_version=3):
    """Retrieve the latest AHRS calibration data for a given CLB

    Parameters
    ----------
    clb_upi: str
    max_version: int, maximum version to check, optional

    Returns
    -------
    Aoff: numpy.array with shape(3,)
    Arot: numpy.array with shape(3,3)
    Hoff: numpy.array with shape(3,)
    Hrot: numpy.array with shape(3,3)

    or None if no calibration found.

    """
    ahrs_upi = km3db.tools.clbupi2compassupi(clb_upi)

    db = km3db.DBManager()

    datasets = []
    for version in range(max_version, 0, -1):
        for n in range(1, 100):
            log.debug("Iteration #{} to get the calib data".format(n))
            url = (
                "show_product_test.htm?upi={0}&"
                "testtype=AHRS-CALIBRATION-v{1}&n={2}&out=xml".format(
                    ahrs_upi, version, n
                )
            )
            log.debug("AHRS calib DB URL: {}".format(url))
            _raw_data = db.get(url).replace("\n", "")
            log.debug("What I got back as AHRS calib: {}".format(_raw_data))
            if len(_raw_data) == 0:
                break
            try:
                xroot = ET.parse(io.StringIO(_raw_data)).getroot()
            except ET.ParseError:
                continue
            else:
                datasets.append(xroot)

    if len(datasets) == 0:
        return None

    latest_dataset = _get_latest_dataset(datasets)
    return _extract_calibration(latest_dataset)


def _get_latest_dataset(datasets):
    """Find the latest valid AHRS calibration dataset"""
    return sorted(datasets, key=lambda d: d.findall(".//EndTime")[0].text)[-1]


def _extract_calibration(xroot):
    """Extract AHRS calibration information from XML root.

    Parameters
    ----------
    xroot: XML root


    Returns
    -------
    Aoff: numpy.array with shape(3,)
    Arot: numpy.array with shape(3,3)
    Hoff: numpy.array with shape(3,)
    Hrot: numpy.array with shape(3,3)

    """
    names = [c.text for c in xroot.findall(".//Name")]
    val = [[i.text for i in c] for c in xroot.findall(".//Values")]

    # The fields has to be reindeced, these are the index mappings
    col_ic = [int(v) for v in val[names.index("AHRS_Matrix_Column")]]
    try:
        row_ic = [int(v) for v in val[names.index("AHRS_Matrix_Row")]]
    except ValueError:
        row_ic = [2, 2, 2, 1, 1, 1, 0, 0, 0]
    try:
        vec_ic = [int(v) for v in val[names.index("AHRS_Vector_Index")]]
    except ValueError:
        vec_ic = [2, 1, 0]

    Aoff_ix = names.index("AHRS_Acceleration_Offset")
    Arot_ix = names.index("AHRS_Acceleration_Rotation")
    Hrot_ix = names.index("AHRS_Magnetic_Rotation")

    Aoff = np.array(val[Aoff_ix])[vec_ic].astype(float)
    Arot = (
        np.array(val[Arot_ix]).reshape(3, 3)[col_ic, row_ic].reshape(3, 3).astype(float)
    )
    Hrot = (
        np.array(val[Hrot_ix]).reshape(3, 3)[col_ic, row_ic].reshape(3, 3).astype(float)
    )

    Hoff = []
    for q in "XYZ":
        values = []
        for t in ("Min", "Max"):
            ix = names.index("AHRS_Magnetic_{}{}".format(q, t))
            values.append(float(val[ix][0]))
        Hoff.append(sum(values) / 2.0)
    Hoff = np.array(Hoff)

    return Aoff, Arot, Hoff, Hrot
