#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Main routine for retrieving reference MeteoSwiss data (e.g. CPC, RZC, POH, etc)
This is meant to be run as a command line command from a slurm script

i.e. ./retrieve_reference_data -t <task_file_name> -c <config_file_name>
- o <output_folder>

IMPORTANT: this function is called by the main routine in database.py
so you should never have to call it manually
--------------
Daniel Wolfensberger, LTE-MeteoSwiss, 2020
Rebecca Gugerli, LTE-MeteoSwiss, 2022
"""

import numpy as np
import pandas as pd
import datetime
import logging
import dask.dataframe as dd
import os

from collections import OrderedDict
from optparse import OptionParser

from rainforest.common import constants
from rainforest.common.lookup import get_lookup
from rainforest.common.utils import read_task_file, envyaml, nested_dict_values
from rainforest.common.retrieve_data import (
    retrieve_prod,
    retrieve_CPCCV,
    retrieve_AQC_XLS,
)
from rainforest.common.io_data import read_cart

logging.basicConfig(level=logging.INFO)

try:
    import pysteps

    _PYSTEPS_AVAILABLE = True
except ImportError:
    _PYSTEPS_AVAILABLE = False

IGNORE_ERRORS = True


class Updater(object):
    def __init__(self, task_file, config_file, output_folder, debug=False):
        """
        Creates an Updater  class instance that allows to add new reference
        data to the database

        Parameters
        ----------
        task_file : str
            The full path to a task file, i.e. a file with the following format
            timestamp, station1, station2, station3...stationN
            These files are generated by the database.py module so normally you
            shouldn't have to create them yourself
        config_file : str
            The full path of a configuration file written in yaml format
            that indicates how the radar retrieval must be done
        output_folder: str
            The full path where the generated files will be stored
        debug: bool
            If set to true will not except any error in the code
        """

        if debug:
            global IGNORE_ERRORS
            IGNORE_ERRORS = False

        self.config = envyaml(config_file)
        self.tasks = read_task_file(task_file)
        self.output_folder = output_folder
        self.downloaded_files = []  # Keeps track of all downloaded files

        self.ref_config = self.config["REFERENCE_RETRIEVAL"]
        self.neighb_x = self.ref_config["NEIGHBOURS_X"]
        self.neighb_y = self.ref_config["NEIGHBOURS_Y"]
        self.products = self.ref_config["PRODUCTS"]
        # Decompose motion vectors variables
        products_decomposed = []
        for prod in self.products:
            if "MV" in prod:
                if not _PYSTEPS_AVAILABLE:
                    logging.error(
                        "Pysteps is not available, product {:s} will not be extracted!"
                    )
                    continue
                products_decomposed.append(prod + "_x")
                products_decomposed.append(prod + "_y")
            else:
                products_decomposed.append(prod)
        self.products = products_decomposed
        self.dims = {
            "np": len(self.products),
            "nnx": len(self.neighb_x),
            "nny": len(self.neighb_y),
        }

    def retrieve_cart_files(self, start_time, end_time, products):
        """
        Retrieves a set of reference product files for a given time range

        Parameters
        ----------
        start_time : datetime.datetime instance
            starting time of the time range
        end_time : datetime.datetime instance
            end time of the time range
        products : list of str
            list of all products to retrieve, must be valid MeteoSwiss product
            names, for example CPC, CPCH, RZC, MZC, BZC, etc
        """

        files_allproducts = {}

        for prod in products:
            try:
                if prod == "CPC" or prod == "CPCH":
                    files = retrieve_prod(
                        start_time,
                        end_time,
                        prod,
                        pattern="*5.801.gif",
                        folder_out=self.config["TMP_FOLDER"],
                    )
                else:
                    files = retrieve_prod(
                        start_time, end_time, prod, folder_out=self.config["TMP_FOLDER"]
                    )

                files_allproducts[prod] = files
                self.downloaded_files.extend(nested_dict_values(files_allproducts))
            except Exception:
                logging.error(
                    """Retrieval for product {:s} at timesteps {:s}-{:s}
                          failed""".format(
                        prod, str(start_time), str(end_time)
                    )
                )
                files_allproducts[prod] = []
                if not IGNORE_ERRORS:
                    raise

        return files_allproducts

    def process_all_timesteps(self):
        """
        Processes all timestaps in the task file
        """

        # Get relevant parameters from user config
        fill_value = self.config["NO_DATA_FILL"]
        nneighb = self.dims["nnx"] * self.dims["nny"]
        logging.info("Products: " + ",".join(self.products))
        logging.info("Nx      : " + ",".join([str(n) for n in self.neighb_x]))
        logging.info("Ny      : " + ",".join([str(n) for n in self.neighb_y]))

        # All 10 min timestamps to process (gauge data)
        all_timesteps = list(self.tasks.keys())

        # LUT to get cartesian data at gauge
        lut_cart = get_lookup("station_to_qpegrid")

        # Initialize output
        data_10minagg = []  # Contains all 10 min data for all products
        data_cst = []  # for time, sta, nx,ny

        if "CPC.CV" in self.products:
            data_cpccv = []  # separate list for cpccv
            data_cst_cpccv = []  # for time, sta, nx,ny
            include_cpccv = True
            self.products.remove("CPC.CV")
            colnames_cpccv = ["TIMESTAMP", "STATION", "NX", "NY"]
            colnames_cpccv.append("CPC.CV")
            # If getting CPC.CV, also get CPC from excel files
            try:
                self.products.remove("CPC_XLS")
            except Exception:
                pass
            data_cpcxls = []
            colnames_cpccv.append("CPC_XLS")
        else:
            include_cpccv = False

        if "AQC_XLS" in self.products:
            data_aqcxls = []
            data_cst_aqcxls = []
            include_aqcxls = True
            self.products.remove("AQC_XLS")
            colnames_aqcxls = ["TIMESTAMP", "STATION", "NX", "NY"]
            colnames_aqcxls.append("AQC_XLS")
        else:
            include_aqcxls = False

        # Initiate
        current_hour = None

        # For motion vectors
        oflow_method = pysteps.motion.get_method(self.ref_config["MV_METHOD"])

        colnames = ["TIMESTAMP", "STATION", "NX", "NY"]
        colnames.extend(self.products)

        for i, tstep in enumerate(all_timesteps):
            logging.info("Processing timestep " + str(tstep))

            # Set t-start -5 minutes to get all the files between, e.g., H:01 and H:10 and log at H:10
            tstart = datetime.datetime.fromtimestamp(
                float(tstep), tz=datetime.timezone.utc
            ) - datetime.timedelta(minutes=8)
            tend = datetime.datetime.fromtimestamp(
                float(tstep), tz=datetime.timezone.utc
            )

            stations_to_get = self.tasks[tstep]

            hour_of_year = datetime.datetime.strftime(tend, "%Y%m%d%H")
            day_of_year = hour_of_year[0:-2]

            current_day = day_of_year

            # Get CPC.CV data
            if hour_of_year != current_hour:
                current_hour = hour_of_year
                if include_cpccv:
                    logging.info(
                        "Retrieving product CPC.CV and CPC_XLS for hour {}".format(
                            current_hour
                        )
                    )

                    # CPC.CV only contains a measurement every full hour
                    tstep_xls = datetime.datetime.strptime(current_hour, "%Y%m%d%H")

                    data_at_stations, data_at_stations_cpc = retrieve_CPCCV(
                        tstep_xls, stations_to_get
                    )

                    # Replace NAN
                    data_at_stations[np.isnan(data_at_stations)] = fill_value
                    data_at_stations_cpc[np.isnan(data_at_stations_cpc)] = fill_value

                    # Assign CPC.CV values to rows corresponding to nx = ny = 0
                    data_cpccv.extend(data_at_stations)
                    data_cpcxls.extend(data_at_stations_cpc)

                    # CPC.CV only contains a measurement every full hour
                    tstep_xls = int(
                        tstep_xls.replace(tzinfo=datetime.timezone.utc).timestamp()
                    )
                    for sta in stations_to_get:
                        data_cst_cpccv.append([tstep_xls, sta, 0, 0])  # nx = ny = 0

                if include_aqcxls:
                    logging.info(
                        "Retrieving product AQC from Excel file for hour {}".format(
                            current_hour
                        )
                    )

                    # AQC_XLS only contains a measurement every full hour
                    tstep_xls = datetime.datetime.strptime(current_hour, "%Y%m%d%H")

                    # Get value
                    data_at_stations = retrieve_AQC_XLS(tstep_xls, stations_to_get)

                    # Replace NAN
                    data_at_stations[np.isnan(data_at_stations)] = fill_value

                    # Assign CPC.CV values to rows corresponding to nx = ny = 0
                    data_aqcxls.extend(data_at_stations)

                    # AQC_XLS only contains a measurement every full hour
                    tstep_xls = int(
                        tstep_xls.replace(tzinfo=datetime.timezone.utc).timestamp()
                    )
                    for sta in stations_to_get:
                        data_cst_aqcxls.append([tstep_xls, sta, 0, 0])  # nx = ny = 0

            # Initialize output
            N, M = len(stations_to_get) * nneighb, self.dims["np"]
            data_allprod = np.zeros((N, M), dtype=np.float32) + np.nan

            # Get data
            baseproducts = [prod for prod in self.products if "MV" not in prod]
            allfiles = self.retrieve_cart_files(tstart, tend, baseproducts)

            for j, prod in enumerate(self.products):
                logging.info("Retrieving product " + prod)
                if "MV" in prod:
                    if "_x" in prod:
                        idx_slice_mv = 0
                        # Motion vector case
                        ###################
                        # Get product for which to compute MV
                        baseprod = prod.strip("MV").split("_")[0]
                        # Initialize output
                        N = len(stations_to_get) * nneighb
                        data_prod = np.zeros((N,), dtype=np.float32) + np.nan

                        try:
                            # For CPC we take only gif
                            files = allfiles[baseprod]

                            R = []
                            for f in files:
                                R.append(read_cart(f))
                            R = np.array(R)
                            R[R < 0] = np.nan
                            mv = oflow_method(R)

                            # Mask mv where there is no rain
                            mask = np.nansum(R, axis=0) <= 0
                            mv[:, mask] = 0
                        except Exception:
                            # fill with missing values, we don't care about the exact dimension
                            mv = np.zeros((2, 1000, 1000)) + fill_value
                            if not IGNORE_ERRORS:
                                raise

                    elif "_y" in prod:  # mv already computed
                        idx_slice_mv = 1

                    idx_row = 0  # Keeps track of the row
                    for sta in stations_to_get:  # Loop on stations
                        for nx in self.neighb_x:
                            for ny in self.neighb_y:
                                strnb = "{:d}{:d}".format(nx, ny)
                                # Get idx of Cart pixel in 2D map
                                idx = lut_cart[sta][strnb]
                                data_prod[idx_row] = mv[idx_slice_mv, idx[0], idx[1]]
                                idx_row += 1

                else:
                    # Normal product case
                    ###################
                    files = allfiles[prod]

                    # # Initialize output
                    N, M = len(stations_to_get) * nneighb, len(files)
                    data_prod = np.zeros((N, M), dtype=np.float32) + np.nan

                    for k, f in enumerate(files):
                        try:
                            proddata = read_cart(f)
                        except Exception:
                            # fill with missing values, we don't care about the exact dimension
                            proddata = np.zeros((1000, 1000)) + np.nan

                        # Threshold radar precip product
                        if prod == "RZC" or prod == "AQC":
                            proddata[proddata < constants.MIN_RZC_VALID] = 0

                        idx_row = 0  # Keeps track of the row
                        for sta in stations_to_get:  # Loop on stations
                            for nx in self.neighb_x:
                                for ny in self.neighb_y:
                                    strnb = "{:d}{:d}".format(nx, ny)
                                    # Get idx of Cart pixel in 2D map
                                    idx = lut_cart[sta][strnb]
                                    data_prod[idx_row, k] = proddata[idx[0], idx[1]]
                                    # Add next row
                                    idx_row += 1

                    data_prod = np.nanmean(data_prod, axis=1)
                    data_prod[np.isnan(data_prod)] = fill_value

                data_allprod[:, j] = data_prod

            for prod in allfiles.keys():
                for f in allfiles[prod]:
                    try:
                        os.remove(f)
                    except Exception:
                        pass

            data_10minagg.extend(data_allprod)

            # Add constant data
            for sta in stations_to_get:
                for nx in self.neighb_x:
                    for ny in self.neighb_y:
                        data_cst.append([tstep, sta, nx, ny])

            # Save data to file if end of loop or new day
            if i == len(all_timesteps) - 1:
                save_output = True
            else:
                next_tstep = datetime.datetime.utcfromtimestamp(all_timesteps[i + 1])
                next_tstep_day = datetime.datetime.strftime(next_tstep, "%Y%m%d%H")[
                    0:-2
                ]
                if current_day != next_tstep_day:
                    save_output = True
                else:
                    save_output = False

            if save_output and len(data_cst):
                logging.info("Saving new table for day {:s}".format(str(current_day)))
                name = self.output_folder + current_day + ".parquet"

                # Check if a file already exists
                file_exists = False
                if os.path.exists(name):
                    file_exists = True

                data_10minagg = np.array(data_10minagg)
                data_cst = np.array(data_cst)
                # Concatenate metadata and product data
                all_data = np.hstack((data_cst, data_10minagg))
                dic = OrderedDict()

                for c, col in enumerate(colnames):
                    data_col = all_data[:, c]
                    isin_listcols = [col in c for c in constants.COL_TYPES.keys()]
                    if any(isin_listcols):
                        idx = np.where(isin_listcols)[0][0]
                        coltype = list(constants.COL_TYPES.values())[idx]
                        try:
                            data_col = data_col.astype(coltype)
                        except Exception:  # for int
                            data_col = data_col.astype(float).astype(coltype)
                            if not IGNORE_ERRORS:
                                raise
                    else:
                        data_col = data_col.astype(np.float32)
                    dic[col] = data_col

                df = pd.DataFrame(dic)

                if include_cpccv:
                    data_cst_cpccv = np.array(data_cst_cpccv)
                    data_cpccv = np.array([data_cpccv]).T
                    data_cpcxls = np.array([data_cpcxls]).T
                    all_data_cpccv = np.hstack(
                        (data_cst_cpccv, data_cpccv, data_cpcxls)
                    )

                    dic_cpccv = OrderedDict()
                    for c, col in enumerate(colnames_cpccv):
                        data_col = all_data_cpccv[:, c]
                        isin_listcols = [col in c for c in constants.COL_TYPES.keys()]
                        if any(isin_listcols):
                            idx = np.where(isin_listcols)[0][0]
                            coltype = list(constants.COL_TYPES.values())[idx]
                            try:
                                data_col = data_col.astype(coltype)
                            except Exception:  # for int
                                data_col = data_col.astype(float).astype(coltype)
                                if not IGNORE_ERRORS:
                                    raise
                        else:
                            data_col = data_col.astype(np.float32)
                        dic_cpccv[col] = data_col

                    dfcpc = pd.DataFrame(dic_cpccv)
                    df = pd.merge(
                        df, dfcpc, on=["STATION", "TIMESTAMP", "NX", "NY"], how="left"
                    )
                    df.replace(np.nan, fill_value, inplace=True)

                # Add AQC Values
                if include_aqcxls:
                    data_cst_aqcxls = np.array(data_cst_aqcxls)
                    data_aqcxls = np.array([data_aqcxls]).T
                    all_data_aqcxls = np.hstack((data_cst_aqcxls, data_aqcxls))

                    dic_aqcxls = OrderedDict()
                    for c, col in enumerate(colnames_aqcxls):
                        data_col = all_data_aqcxls[:, c]
                        isin_listcols = [col in c for c in constants.COL_TYPES.keys()]
                        if any(isin_listcols):
                            idx = np.where(isin_listcols)[0][0]
                            coltype = list(constants.COL_TYPES.values())[idx]
                            try:
                                data_col = data_col.astype(coltype)
                            except Exception:  # for int
                                data_col = data_col.astype(float).astype(coltype)
                                if not IGNORE_ERRORS:
                                    raise
                        else:
                            data_col = data_col.astype(np.float32)
                        dic_aqcxls[col] = data_col

                    dfcpc = pd.DataFrame(dic_aqcxls)
                    df = pd.merge(
                        df, dfcpc, on=["STATION", "TIMESTAMP", "NX", "NY"], how="left"
                    )
                    df.replace(np.nan, fill_value, inplace=True)

                if not file_exists:
                    logging.info("Saving file " + name)
                    df.to_parquet(name, compression="gzip", index=False)
                else:
                    logging.info(
                        "Saving file: "
                        + name
                        + " as it already exists, I will append to it"
                    )
                    # Rename old file to be able to delete it afterwards
                    name_old = name[0:-8] + "_old" + name[-8::]
                    os.rename(name, name_old)
                    # Read dask DataFrame and convert it to Pandas DataFrame
                    df_old = dd.read_parquet(name_old).compute()
                    # Merge the old and new one and drop duplicate rows
                    df_join = df_old._append(df).drop_duplicates()
                    # Save the new file and delete the old one
                    df_join.to_parquet(name, compression="gzip", index=False)
                    os.remove(name_old)

                # Reset lists
                data_10minagg = []  # separate list for cpccv
                data_cst = []  # for time, sta, nx,ny
                if include_cpccv:
                    data_cst_cpccv = []  # for time, sta, nx,ny
                    data_cpccv = []  # separate list for cpccv
                    data_cpcxls = []  # separate list for cpc from xls files
                if include_aqcxls:
                    data_cst_aqcxls = []
                    data_aqcxls = []

    def final_cleanup(self):
        """
        Performs a final cleanup by checking if all files in downloaded_files
        deleted
        """
        for f in self.downloaded_files:
            if os.path.exists(f):
                try:
                    os.remove(f)
                except PermissionError as e:
                    logging.error(e)
                    logging.error("Could not delete file {:s}".format(f))


if __name__ == "__main__":
    parser = OptionParser()

    parser.add_option(
        "-c",
        "--configfile",
        dest="config_file",
        help="Specify the user configuration file to use",
        metavar="CONFIG",
    )

    parser.add_option(
        "-t",
        "--taskfile",
        dest="task_file",
        default=None,
        help="Specify the task file to process",
        metavar="TASK",
    )

    parser.add_option(
        "-o",
        "--output",
        dest="output_folder",
        default="/tmp/",
        help="Specify the output directory",
        metavar="FOLDER",
    )

    (options, args) = parser.parse_args()

    u = Updater(options.task_file, options.config_file, options.output_folder)
    u.process_all_timesteps()
    u.final_cleanup()
