#!/usr/bin/env python

# Copyright (c) 2015-2023 Mikael Lund
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# THIS SOFTWARE IS PROVIDED ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL Berkeley Software Design, Inc. BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#
# Module use:
#
#   from loadgrace import LoadGrace
#   grace = LoadGrace("filename.agr")
#   for (x, y), label, comment in grace:
#       plt.plot(x, y, label=label)
#
# Command line use:
#
#   $ loadgrace.py --plot filename.agr
#       ┌─────────────────────────────────────────────────────────────────────────────────────┐
#   23.5┤ ▞▞ mmAm                  ▌                                                          │
#       │ ▞▞ mmAm (ideal)          ▐                                                          │
#       │ ▞▞ hisactophilin          ▌                                                         │
#       │ ▞▞ hisactophilin (ideal)  ▐                                                         │
#   19.6┤                       ▌    ▌                                                        │
#       │                      ▐     ▚                                                        │
#       │                      ▞     ▐                                                        │
#       │                     ▗▘     ▐                                                        │
#   15.7┤                     ▐       ▌                                                       │
#       │                     ▌       ▌                                                       │
#       │                     ▌       ▐                                                       │
#       │                    ▐        ▐                               ▄▄▄                     │
#   11.8┤                    ▐         ▌                             ▞   ▚      ▄▄            │
#       │                    ▌         ▝▖                          ▗▀     ▚▖  ▗▞  ▚           │
#       │                    ▌          ▌                          ▞       ▝▄▄▘   ▝▖          │
#       │                   ▐           ▐                         ▐                ▝▖         │
#       │                   ▞           ▝▖                       ▗▘                 ▚         │
#    7.8┤      ▄▄▄▄▄▄▚▄▄▖  ▗▘            ▚   ▄▞▀▖                ▌                  ▝▖▄▞▄▄▄   │
#       │    ▄▀         ▝▀▀▞▄▖           ▝▖ ▐   ▝▚              ▐               ▗▞▀▀▘▌     ▀▀▀│
#       │ ▗▞▀             ▗▘ ▝▀▚▄▖        ▝▗▘     ▌            ▗▘             ▄▞▘    ▝▖       │
#       │▀▘               ▐      ▝▀▀▚▄▄▖  ▗▘▄▖    ▝▄           ▞           ▗▄▀        ▝▖      │
#    3.9┤                ▗▘       ▗▄   ▝▀▄▘▖ ▝▀▚▖  ▝▖         ▞         ▗▄▀▘           ▚      │
#       │               ▗▞     ▗▞▀▘▄▀▀▀▀▀▀▀▀▀▚▄▄▖▚▖ ▝▖       ▞       ▗▄▄▄               ▌     │
#       │    ▄▄▚▄▄▄▄▄▄▄▄▄▄▄▄▄▗▞▘▀               ▝▀▀▀▄▝▄▖▄▄▄▗▞▘▄▄▄▄▗▄▀▀   ▀▄▄            ▝▄▖   │
#       │▄▄▀▀        ▗▄▘  ▗▄▞▘                       ▀▀▝▚▄▖▄▄▄▄▄▞▀▘▄▄▄▄▄▄▄▄▞▀▄▄▖▀▀▄▄▄▄▄▄  ▝▄  │
#    0.0┤▄▄▄▄▄▄▄▄▄▄▄▄▄▄▞▀▀▘                               ▝▀▀▀                 ▝▀▀▀▀▀▄▄▄▄▄▄▄▄▄│
#       └┬────────────────────┬────────────────────┬────────────────────┬────────────────────┬┘
#       0.0                  3.5                  7.0                 10.5                14.0
#

import numpy as np
import re
import sys


class LoadGrace:
    """
    Simple module to load a Grace file (.agr) plot and extract
    legends, comments, and extract data into numpy arrays.
    """

    def __init__(self, filename):
        self._sets = []
        self._legends = []
        self._comments = []

        with open(filename) as file:
            for line in file:
                if re.compile(r"@\s+s(.*) legend", re.IGNORECASE).search(line):
                    self._legends.append(line.split('"')[1])

                if re.compile(r"@\s+s(.*) comment", re.IGNORECASE).search(line):
                    self._comments.append(line.split('"')[1])

                if "@target" in line:
                    values = []
                    next(file)
                    for line in file:
                        if line != "&\n":
                            values.append(np.fromstring(line, sep=" ", count=2))
                        else:
                            self._sets.append(np.array(values))
                            break

    def sets(self):
        """Data sets"""
        return self._sets

    def legends(self):
        """Legends for each dataset"""
        return self._legends

    def comments(self):
        """Comments for each dataset"""
        return self._comments

    def __len__(self):
        """Number of datasets"""
        return len(self._sets)

    def __getitem__(self, key):
        """Get all data by index and enable iteration
        Returns x, y, legend, comment
        """
        x = self._sets[key][:, 0]
        y = self._sets[key][:, 1]
        return (x, y), self._legends[key], self._comments[key]

    def plot(self):
        """Plot data using plotext"""
        try:
            import plotext as plt
        except ImportError:
            print(
                "Install plotext with e.g. `pip install plotext` to plot data in the terminal"
            )
            sys.exit(1)

        for dataset, legend, comment in self:
            plt.plot(dataset[0], dataset[1], label=legend)

        plt.show()


# If run as main program
if __name__ == "__main__":
    import argparse

    ps = argparse.ArgumentParser(description="Extract data from xmgrace (.agr) files")
    ps.add_argument("filename", type=str, help="xmgrace file")
    ps.add_argument("-v", "--version", action="version", version="%(prog)s 0.1.0")
    ps.add_argument(
        "-p",
        "--plot",
        action="store_true",
        help="Plot data in terminal (requires plotext)",
    )
    args = ps.parse_args()

    grace = LoadGrace(args.filename)

    if args.plot:
        grace.plot()
    else:
        for dataset, legend, comment in grace:
            print(f"# Legend: {legend} Comment: {comment}")
            print(dataset[0], dataset[1])
            print()

    sys.exit(0)