#!/usr/bin/env seiscomp-python

############################################################################
# Copyright (C) 2016 by gempa GmbH                                         #
#                                                                          #
# All Rights Reserved.                                                     #
#                                                                          #
# NOTICE: All information contained herein is, and remains                 #
# the property of gempa GmbH and its suppliers, if any. The intellectual   #
# and technical concepts contained herein are proprietary to gempa GmbH    #
# and its suppliers.                                                       #
# Dissemination of this information or reproduction of this material       #
# is strictly forbidden unless prior written permission is obtained        #
# from gempa GmbH.                                                         #
#                                                                          #
#  Author: Enrico Ellguth, Dirk Roessler                                   #
#  Email: enrico.ellguth@gempa.de, roessler@gempa.de                       #
############################################################################

import datetime
import os
import sys

from seiscomp import core, datamodel, io
from seiscomp.client import Application
from seiscomp import geo


def str2time(timestring):
    """
    Liberally accept many time string formats and convert them to a
    seiscomp.core.Time
    """

    timestring = timestring.strip()
    for c in ["-", "/", ":", "T", "Z"]:
        timestring = timestring.replace(c, " ")
    timestring = timestring.split()
    assert 3 <= len(timestring) <= 6
    timestring.extend((6 - len(timestring)) * ["0"])
    timestring = " ".join(timestring)

    fmt = "%Y %m %d %H %M %S"
    if timestring.find(".") != -1:
        fmt += ".%f"

    t = core.Time()
    t.fromString(timestring, fmt)
    return t


def utc():
    return datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)


class DumpPicks(Application):
    def __init__(self, argc, argv):
        Application.__init__(self, argc, argv)
        self.output = "-"
        self.type = "0"
        self.margin = [300]
        self.originID = None
        self.bbox = None
        self.noamp = False
        self.automatic = False
        self.manual = False
        self.checkInventory = False
        self.author = None
        self.hours = None
        self.minutes = None
        self.start = None
        self.end = None
        self.network = None
        self.station = None
        self.tmin = str2time("1970-01-01 00:00:00")
        self.tmax = str2time(str(utc()))
        self.delay = None

        self.setMessagingEnabled(False)
        self.setDatabaseEnabled(True, True)

    def createCommandLineDescription(self):
        self.commandline().addGroup("Dump")

        self.commandline().addStringOption(
            "Dump",
            "hours",
            "Start search hours before now considering object time, not creation time. "
            "If --minutes is given as well they will be added. "
            "If set, --time-window, --start, --end are ignored.",
        )

        self.commandline().addStringOption(
            "Dump",
            "minutes",
            "Start search minutes before now considering object time, not creation time. "
            "If --hours is given as well they will be added. "
            "If set, --time-window, --start, --end are ignored.",
        )

        self.commandline().addStringOption(
            "Dump",
            "start",
            "Start time of search until now considering object time, not creation time."
            " If set, --time-window is ignored.",
        )

        self.commandline().addStringOption(
            "Dump",
            "end",
            "End time of search considering object time, not creation time. If set, "
            "--time-window is ignored.",
        )

        self.commandline().addStringOption(
            "Dump",
            "time-window,t",
            "Specify time window to search picks and amplitudes by their time. Use one "
            "single string which must be enclosed by quotes in case of spaces in the "
            "time string. Times are of course in UTC and separated by a tilde '~'. "
            "Uses: 1970-01-01 00:00:00 to now if not set.",
        )

        self.commandline().addStringOption(
            "Dump",
            "maximum-delay",
            "Maximum allowed delay of picks or amplitudes, hence the difference between"
            " creation time and actual time value. Allows identifcation of picks found "
            "in real time.",
        )

        self.commandline().addStringOption(
            "Dump",
            "region,r",
            "Dump picks only from sensors in given region. Implies loading an "
            "inventory.\n"
            "Format: minLat,minLon,maxLat,maxLon \n"
            "Default: -90,-180,90,180 if not set.",
        )

        self.commandline().addOption(
            "Dump",
            "check-inventory,c",
            "Dump picks only when corresponding streams are found in inventory.",
        )

        self.commandline().addStringOption(
            "Dump",
            "origin,O",
            "Origin ID. Dump all "
            "picks associated with the origin that has the given origin ID.",
        )

        self.commandline().addOption("Dump", "manual,m", "Dump only manual picks.")

        self.commandline().addOption(
            "Dump", "automatic,a", "Dump only automatic picks."
        )

        self.commandline().addOption(
            "Dump",
            "no-amp,n",
            "Do not dump amplitudes from picks. "
            "Amplitudes are not required by scanloc.",
        )

        self.commandline().addStringOption(
            "Dump", "author", "Filter picks by the given author."
        )

        self.commandline().addStringOption(
            "Dump",
            "net-sta",
            "Filter picks and amplitudes by given network code or "
            "network and station code. Format: NET or NET.STA.",
        )

        self.commandline().addGroup("Output")
        self.commandline().addStringOption(
            "Output",
            "output,o",
            "Name of output file. If not given, all data is written to stdout.",
        )
        self.commandline().addStringOption(
            "Output",
            "type",
            f"Type of output format. Default: {self.type}.\n"
            "0 / scml: SCML containing all objects (default if option is not used)\n"
            "1 / streams: Time windows and streams for all picks like in scevtstreams\n"
            "2 / caps: Time windows and streams in capstool format\n"
            "3 / fdsnws: Time windows and streams in FDSN dataselect webservice POST \
format\n"
            "Except for type 0, only picks are considered ignoring all other objects.",
        )
        self.commandline().addOption(
            "Output",
            "formatted,f",
            "Output formatted XML. Default is unformatted. Applies only for type 0.",
        )
        self.commandline().addStringOption(
            "Output",
            "margin",
            "Time margin applied around pick times along with --type = [1:]. Use 2 "
            "comma-separted values (before,after) for asymmetric margins, e.g. "
            f"--margin 120,300. Default: {self.margin[0]} s.",
        )

    def printUsage(self):
        print(
            f"""Usage:
  {os.path.basename(__file__)} [options]

Read picks and amplitudes from database and dump them to a file or to standard output.\
"""
        )

        Application.printUsage(self)

        print(
            f"""Examples:
Dump all picks within a region and a period of time
  {os.path.basename(__file__)} -d localhost -t 2023-01-20T13:52:00~2023-01-20T13:57:00\
 -r "-10,-90,10,120"

Search 24 hours before now for automatic picks from author "scautopick" with low delay \
ignoring amplitudes
  {os.path.basename(__file__)} -d localhost --hours 24 -a -n --author "scautopick" \
--maximum-delay 60

Dump the streams of picks with time windows fetching the corresponding data from a \
local CAPS server
  {os.path.basename(__file__)} -d localhost --type 2 --margin 60 | capstool \
-H localhost -o data.mseed

Dump the streams of picks with time windows fetching the corresponding data from a \
local SDS archive
  {os.path.basename(__file__)} -d localhost --type 1 --margin 60 | scart -dsE -l - \
/archive -o data.mseed
"""
        )

    def init(self):
        if not Application.init(self):
            return False

        try:
            self.output = self.commandline().optionString("output")
        except RuntimeError:
            pass

        try:
            self.type = self.commandline().optionString("type")
        except RuntimeError:
            pass

        if self.type == "scml":
            self.type = "0"
        elif self.type == "streams":
            self.type = "1"
        elif self.type == "caps":
            self.type = "2"
        elif self.type == "fdsnws":
            self.type = "3"

        try:
            self.margin = self.commandline().optionString("margin").split(",")
        except RuntimeError:
            pass

        try:
            self.originID = self.commandline().optionString("origin")
        except RuntimeError:
            pass

        if not self.originID:
            try:
                boundingBox = self.commandline().optionString("region")
                self.bbox = boundingBox.split(",")
                if len(self.bbox) != 4:
                    print(
                        "Invalid region given, expected lat0,lon0,lat1,lon1",
                        file=sys.stderr,
                    )
                    return False

                self.bbox[0] = str(geo.GeoCoordinate.normalizeLat(float(self.bbox[0])))
                self.bbox[1] = str(geo.GeoCoordinate.normalizeLon(float(self.bbox[1])))
                self.bbox[2] = str(geo.GeoCoordinate.normalizeLat(float(self.bbox[2])))
                self.bbox[3] = str(geo.GeoCoordinate.normalizeLon(float(self.bbox[3])))

                self.checkInventory = True
            except RuntimeError:
                boundingBox = "-90,-180,90,180"
                self.bbox = boundingBox.split(",")

            print("Settings", file=sys.stderr)
            print(
                f"  + considered region: {self.bbox[0]} - {self.bbox[2]} deg North, "
                f"{self.bbox[1]} - {self.bbox[3]} deg East",
                file=sys.stderr,
            )

            try:
                self.hours = float(self.commandline().optionString("hours"))
            except RuntimeError:
                pass

            try:
                self.minutes = float(self.commandline().optionString("minutes"))
            except RuntimeError:
                pass

            try:
                self.start = self.commandline().optionString("start")
            except RuntimeError:
                pass

            try:
                self.end = self.commandline().optionString("end")
            except RuntimeError:
                pass

            delta = 0.0
            if self.hours:
                delta = self.hours * 60
            if self.minutes:
                delta += self.minutes

            if self.hours or self.minutes:
                print(
                    "  + time window set by hours and/or minutes option: ignoring all "
                    "other time parameters",
                    file=sys.stderr,
                )
                dt = datetime.timedelta(minutes=delta)

                self.tmin = str2time(str(utc() - dt))
                self.tmax = str2time(str(utc()))
                self.start = None
                self.end = None

            else:
                if self.start:
                    print(
                        "  + time window set by start option: ignoring --time-window",
                        file=sys.stderr,
                    )
                    self.tmin = str2time(self.start)

                if self.end:
                    print(
                        "  + time window set by end option: ignoring --time-window",
                        file=sys.stderr,
                    )
                    self.tmax = str2time(self.end)

                if not self.start and not self.end:
                    try:
                        self.tmin, self.tmax = map(
                            str2time,
                            self.commandline().optionString("time-window").split("~"),
                        )
                        print(
                            "  + time window set by time-window option", file=sys.stderr
                        )
                    except RuntimeError:
                        print(
                            "  + no time window given exlicitly: Assuming defaults",
                            file=sys.stderr,
                        )

            print(
                f"  + considered time window: {str(self.tmin)} - {str(self.tmax)}",
                file=sys.stderr,
            )
        else:
            print(
                "  + searching for picks is based on originID, ignoring "
                "region and time window",
                file=sys.stderr,
            )

        try:
            self.delay = float(self.commandline().optionString("maximum-delay"))
        except RuntimeError:
            pass

        if not self.checkInventory:
            self.checkInventory = self.commandline().hasOption("check-inventory")

        if self.checkInventory:
            print(
                "  + dumping only picks for streams found in inventory", file=sys.stderr
            )
        else:
            print(
                "  + do not consider inventory information for dumping picks",
                file=sys.stderr,
            )

        if self.commandline().hasOption("no-amp"):
            self.noamp = True
        else:
            self.noamp = False

        if self.type != "0":
            self.noamp = True

        if self.noamp:
            print("  + dumping picks without amplitudes", file=sys.stderr)
        else:
            print("  + dumping picks with amplitudes", file=sys.stderr)

        if self.commandline().hasOption("manual"):
            self.manual = True
            print("  + dumping only manual objects", file=sys.stderr)
        else:
            self.manual = False
            print("  + considering also manual objects", file=sys.stderr)

        if self.commandline().hasOption("automatic"):
            if not self.manual:
                self.automatic = True
                print("  + dumping only automatic picks", file=sys.stderr)
            else:
                print(
                    "EXIT - Script was started with competing options -a and -m",
                    file=sys.stderr,
                )
                return False
        else:
            self.automatic = False
            print("  + considering also automatic objects", file=sys.stderr)

        try:
            self.author = self.commandline().optionString("author")
        except RuntimeError:
            pass

        networkStation = None
        try:
            networkStation = self.commandline().optionString("net-sta")
            print(
                f"  + filter objects by network / station code: {networkStation}",
                file=sys.stderr,
            )
        except RuntimeError:
            pass

        if networkStation:
            try:
                self.network = networkStation.split(".")[0]
            except IndexError:
                print(
                    f"Error in network code '{networkStation}': Use '--net-sta' with "
                    "format NET or NET.STA",
                    file=sys.stderr,
                )
                return False

            try:
                self.station = networkStation.split(".")[1]
            except IndexError:
                print(
                    f"  + no station code given in '--net-sta {networkStation}' - "
                    "using all stations from network",
                    file=sys.stderr,
                )

        return True

    def run(self):
        db = self.database()

        def _T(name):
            return db.convertColumnName(name)

        def _time(time):
            return db.timeToString(time)

        colLat, colLon = _T("latitude"), _T("longitude")

        dbq = self.query()
        ep = datamodel.EventParameters()
        picks = []
        noAmps = 0

        if self.originID:
            for p in dbq.getPicks(self.originID):
                picks.append(datamodel.Pick.Cast(p))

            for p in picks:
                dbq.loadComments(p)
                ep.add(p)

            if not self.noamp:
                for a in dbq.getAmplitudesForOrigin(self.originID):
                    amp = datamodel.Amplitude.Cast(a)
                    ep.add(amp)

        else:
            fmt = "%Y-%m-%d %H:%M:%S"
            if self.checkInventory:
                q = (
                    "select distinct(PPick.%s), Pick.* "
                    "from PublicObject as PPick, Pick, Network, Station, SensorLocation "
                    "where PPick._oid=Pick._oid and Network._oid=Station._parent_oid and "
                    "Station._oid=SensorLocation._parent_oid and Station.%s >= %s and "
                    "Station.%s <= %s and Station.%s >= %s and Station.%s <= %s and "
                    "SensorLocation.%s=Pick.%s and SensorLocation.%s <= Pick.%s and "
                    "(SensorLocation.%s is null or SensorLocation.%s > Pick.%s) and "
                    "Station.%s=Pick.%s and Network.%s=Pick.%s and "
                    "Pick.%s >= '%s' and Pick.%s < '%s'"
                    ""
                    % (
                        _T("publicID"),
                        colLat,
                        self.bbox[0],
                        colLat,
                        self.bbox[2],
                        colLon,
                        self.bbox[1],
                        colLon,
                        self.bbox[3],
                        _T("code"),
                        _T("waveformID_locationCode"),
                        _T("start"),
                        _T("time_value"),
                        _T("end"),
                        _T("end"),
                        _T("time_value"),
                        _T("code"),
                        _T("waveformID_stationCode"),
                        _T("code"),
                        _T("waveformID_networkCode"),
                        _T("time_value"),
                        self.tmin.toString(fmt),
                        _T("time_value"),
                        self.tmax.toString(fmt),
                    )
                )
            else:
                q = (
                    "select distinct(PPick.%s), Pick.* "
                    "from PublicObject as PPick, Pick "
                    "where PPick._oid=Pick._oid and "
                    "Pick.%s >= '%s' and Pick.%s < '%s'"
                    ""
                    % (
                        _T("publicID"),
                        _T("time_value"),
                        self.tmin.toString(fmt),
                        _T("time_value"),
                        self.tmax.toString(fmt),
                    )
                )

            if self.manual:
                q = q + f" and Pick.{_T('evaluationMode')} = 'manual' "

            if self.automatic:
                q = q + f" and Pick.{_T('evaluationMode')} = 'automatic' "

            if self.author:
                q = q + f" and Pick.{_T('creationInfo_author')} = '{self.author}' "

            if self.network:
                q = q + f" and Pick.{_T('waveformID_networkCode')} = '{self.network}' "

            if self.station:
                q = q + f" and Pick.{_T('waveformID_stationCode')} = '{self.station}' "

            for p in dbq.getObjectIterator(q, datamodel.Pick.TypeInfo()):
                pick = datamodel.Pick.Cast(p)
                if (
                    self.delay
                    and float(pick.creationInfo().creationTime() - pick.time().value())
                    > self.delay
                ):
                    continue
                picks.append(pick)

            for p in picks:
                dbq.loadComments(p)
                ep.add(p)

            if not self.noamp:
                if self.checkInventory:
                    q = (
                        "select distinct(PAmplitude.%s), Amplitude.* "
                        "from PublicObject as PAmplitude, Amplitude, PublicObject \
                        as PPick, Pick, Network, Station, SensorLocation "
                        "where PAmplitude._oid=Amplitude._oid and "
                        "PPick._oid=Pick._oid and Network._oid=Station._parent_oid and "
                        "Station._oid=SensorLocation._parent_oid and Station.%s >= %s and "
                        "Station.%s <= %s and Station.%s >= %s and Station.%s <= %s and "
                        "SensorLocation.%s=Pick.%s and SensorLocation.%s <= Pick.%s and "
                        "(SensorLocation.%s is null or SensorLocation.%s > Pick.%s) and "
                        "Station.%s=Pick.%s and Network.%s=Pick.%s and "
                        "Pick.%s >= '%s' and Pick.%s < '%s' and PPick.%s=Amplitude.%s"
                        ""
                        % (
                            _T("publicID"),
                            colLat,
                            self.bbox[0],
                            colLat,
                            self.bbox[2],
                            colLon,
                            self.bbox[1],
                            colLon,
                            self.bbox[3],
                            _T("code"),
                            _T("waveformID_locationCode"),
                            _T("start"),
                            _T("time_value"),
                            _T("end"),
                            _T("end"),
                            _T("time_value"),
                            _T("code"),
                            _T("waveformID_stationCode"),
                            _T("code"),
                            _T("waveformID_networkCode"),
                            _T("time_value"),
                            self.tmin.toString(fmt),
                            _T("time_value"),
                            self.tmax.toString(fmt),
                            _T("publicID"),
                            _T("pickID"),
                        )
                    )
                else:
                    q = (
                        "select distinct(PAmplitude.%s), Amplitude.* "
                        "from PublicObject as PAmplitude, Amplitude, PublicObject as PPick, Pick "
                        "where PAmplitude._oid=Amplitude._oid and PPick._oid=Pick._oid and "
                        "Pick.%s >= '%s' and Pick.%s < '%s' and PPick.%s=Amplitude.%s"
                        ""
                        % (
                            _T("publicID"),
                            _T("time_value"),
                            self.tmin.toString(fmt),
                            _T("time_value"),
                            self.tmax.toString(fmt),
                            _T("publicID"),
                            _T("pickID"),
                        )
                    )

                if self.manual:
                    q = q + f" and Pick.{_T('evaluationMode')} = 'manual' "
                if self.automatic:
                    q = q + f" and Pick.{_T('evaluationMode')} = 'automatic' "

                if self.author:
                    q = q + f" and Pick.{_T('creationInfo_author')} = '{self.author}' "

                if self.network:
                    q = q + " and Pick.%s = '%s' " % (
                        _T("waveformID_networkCode"),
                        self.network,
                    )

                if self.station:
                    q = q + " and Pick.%s = '%s' " % (
                        _T("waveformID_stationCode"),
                        self.station,
                    )

                for a in dbq.getObjectIterator(q, datamodel.Amplitude.TypeInfo()):
                    amp = datamodel.Amplitude.Cast(a)
                    if (
                        self.delay
                        and float(
                            amp.creationInfo().creationTime()
                            - amp.timeWindow().reference()
                        )
                        > self.delay
                    ):
                        continue
                    ep.add(amp)
                    noAmps += 1

        if self.type == "0":
            ar = io.XMLArchive()
            ar.create(self.output)
            ar.setFormattedOutput(self.commandline().hasOption("formatted"))
            ar.writeObject(ep)
            ar.close()
        elif self.type in ["1", "2", "3"]:
            if len(picks) == 0:
                print(
                    "No picks are found and written",
                    file=sys.stderr,
                )
                return False

            # convert times to string depending on requested output format
            # time and line format
            if self.type == "2":
                timeFMT = "%Y,%m,%d,%H,%M,%S"
                lineFMT = "{0} {1} {2} {3} {4} {5}"
            elif self.type == "3":
                timeFMT = "%FT%T"
                lineFMT = "{2} {3} {4} {5} {0} {1}"
            else:
                timeFMT = "%F %T"
                lineFMT = "{0};{1};{2}.{3}.{4}.{5}"

            lines = set()
            for pick in picks:
                net = pick.waveformID().networkCode()
                station = pick.waveformID().stationCode()
                loc = pick.waveformID().locationCode()
                channelGroup = f"{pick.waveformID().channelCode()[:2]}*"

                # FDSNWS requires empty location to be encoded by 2 dashes
                if not loc and self.type == "3":
                    loc = "--"

                # add some marging to picks times
                minTime = pick.time().value() - core.TimeSpan(float(self.margin[0]))
                maxTime = pick.time().value() + core.TimeSpan(float(self.margin[-1]))
                minTime = minTime.toString(timeFMT)
                maxTime = maxTime.toString(timeFMT)
                lines.add(
                    lineFMT.format(minTime, maxTime, net, station, loc, channelGroup)
                )

            if self.output == "-":
                out = sys.stdout
            else:
                print(f"Output data to file: {self.output}", file=sys.stderr)
                try:
                    out = open(self.output, "w", encoding="utf8")
                except Exception:
                    print("Cannot create output file '{self.output}'", file=sys.stderr)
                    return False

            for line in sorted(lines):
                print(line, file=out)

            if self.output != "-":
                out.close()
        else:
            print(
                f"Unspupported output format '{self.type}': No objects are written",
                file=sys.stderr,
            )
            return False

        print(
            f"Saved: {len(picks):d} picks, {noAmps:d} amplitudes",
            file=sys.stderr,
        )
        return True


def main(argv):
    app = DumpPicks(len(argv), argv)
    return app()


if __name__ == "__main__":
    sys.exit(main(sys.argv))
