#!/usr/bin/env python
"""SlitherStats computes statistics from the results of HiSlither."""

# Copyright 2004-2020, Arizona Board of Regents on behalf of the Lunar and
# Planetary Laboratory at the University of Arizona.
#   - Orignal Perl and Octave programs.
#
# Copyright 2020, Ross A. Beyer (rbeyer@seti.org)
#   - Elements of this Python program are are based on the original Perl
#     and Octave but the logic here is rewritten from scratch to emulate
#     functionality.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# This program is based on HiColor version 5.4.5 (2021/03/02),
# the SlitherStats Perl program ($Revision: 1.9 $ $Date: 2021/03/02 20:21:50 $)
# and on the Octave program
# readSlitherPolyfit_octave.m ($Id: v 1.6 2021/03/02 21:35:38)
# by Sarah Sutton and Audrie Fennema as employees of the University of Arizona.

import argparse
import csv
import logging
import os

import matplotlib.pyplot as plt
import numpy as np
import numpy.polynomial.polynomial as P

import hiproc.util as util

logger = logging.getLogger(__name__)


def arg_parser():
    parser = argparse.ArgumentParser(
        description=__doc__, parents=[util.parent_parser()]
    )
    parser.add_argument(
        "-p", "--plot", action="store_true", help="Display plots"
    )
    parser.add_argument("slithertxts", metavar="slither.txt files", nargs="+")
    return parser


def main():
    args = arg_parser().parse_args()

    util.set_logger(args.verbose, args.logfile, args.log)

    for s in args.slithertxts:
        slitherstats = Polyfit(s, args.plot)
        print(f"Stats for {s}")
        print(
            "Max Diff: {:.2f}\nAve Diff: {:.2f}\nStd Dev: {:.2f}".format(
                *slitherstats
            )
        )


def read_slithertxt(filename: os.PathLike) -> tuple:
    """Reads the slither.txt file and returns these three columns as numpy
    arrays in a tuple: MatchLine, LineOffset, and SampOffset.
    """

    reg_statistics = list()
    consume = False
    with open(filename, "r") as f:
        for line in f:
            if not consume:
                if line.startswith("  FromLine"):
                    reg_statistics.append(line)
                    consume = True
                else:
                    continue
            else:
                if line.isspace():
                    consume = False
                    break
                else:
                    reg_statistics.append(line)

    matchline = list()
    lineoffset = list()
    sampoffset = list()

    dialect = csv.Dialect
    dialect.delimiter = " "
    dialect.skipinitialspace = True
    dialect.quoting = csv.QUOTE_NONE
    dialect.lineterminator = "\n"

    reader = csv.DictReader(reg_statistics, dialect=dialect)
    for row in reader:
        matchline.append(float(row["MatchLine"]))
        lineoffset.append(float(row["LineOffset"]))
        sampoffset.append(float(row["SampOffset"]))

    return (np.array(matchline), np.array(lineoffset), np.array(sampoffset))


def plot_stats(
    filename,
    xline,
    magnitudes,
    poly_values,
    poly_deg,
    absdiff,
    maxdiff,
    avediff,
):

    plt.ioff()
    fig, (ax0, ax1) = plt.subplots(2, 1, sharex=True)
    fig.suptitle(f"SliterStats for {filename}")

    ax0.set_ylabel("Slither Magnitude (pixels)")
    ax0.scatter(xline, magnitudes, label="Magnitude")
    ax0.scatter(xline, poly_values, label=f"Polynomial Fit, degree {poly_deg}")
    ax0.legend()

    ax1.set_ylabel("Jitter Amplitude (pixels)")
    ax1.set_xlabel("Image Line")
    ax1.scatter(xline, absdiff)
    ax1.axhline(y=avediff)

    ax1.text(
        1,
        1,
        "Maximum Diff: {:.2f}\nAverage Diff: {:.2f}".format(maxdiff, avediff),
        verticalalignment="top",
        horizontalalignment="right",
        transform=ax1.transAxes,
    )

    plt.show()

    return


def Polyfit(slither_file: os.PathLike, plot=False) -> tuple:
    """Read text files generated by slither.
    Calculate the error magnitude from the control points.
    The local minima and maxima of the plotted points are found.
    These points are interpolated, to create 2 functions.
    Statistics about the difference between these two fuctions are generated
    to give information about the jitter conditions for that observation.
    """

    (xline, lineoffsets, sampoffsets) = read_slithertxt(slither_file)

    magnitudes = np.hypot(lineoffsets, sampoffsets)

    poly_deg = 6
    if np.size(xline) > 40000:
        poly_deg = 12

    coefficients = P.polyfit(xline, magnitudes, poly_deg)
    poly_values = P.polyval(xline, coefficients)

    diff = magnitudes - poly_values

    absdiff = np.absolute(diff)

    maxdiff = np.amax(absdiff)
    avediff = np.mean(absdiff)
    stddiff = np.std(absdiff)

    if plot:
        plot_stats(
            slither_file,
            xline,
            magnitudes,
            poly_values,
            poly_deg,
            absdiff,
            maxdiff,
            avediff,
        )

    return (maxdiff, avediff, stddiff)
