#!/usr/bin/env python
"""
Runs an ILAMB study.
"""
import mpi4py.rc

mpi4py.rc.threads = False
import logging

import yaml
from ILAMB.ModelResult import ModelResult

try:
    from ILAMB.point_result import ModelPointResult
except:
    ModelPointResult = None
try:
    from ILAMB.e3sm_result import E3SMResult
except:
    E3SMResult = None
import argparse
import datetime
import glob
import inspect
import os
import pickle
import re
import sys
import time
from traceback import format_exc

import matplotlib.colors as clr
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from ILAMB import ilamblib as il
from ILAMB.Post import RegisterCustomColormaps
from ILAMB.Regions import Regions
from ILAMB.Scoreboard import Scoreboard
from mpi4py import MPI
from netCDF4 import Dataset

if "wetdry" not in plt.colormaps():
    RegisterCustomColormaps()
import platform

# MPI stuff
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
proc = np.zeros(size)
ierr = np.zeros(size)

# Some color constants for printing to the terminal
OK = "\033[92m"
FAIL = "\033[91m"
ENDC = "\033[0m"


def InitializeModels(
    model_root,
    models=[],
    verbose=False,
    filter="",
    regex="",
    model_year=[],
    log=True,
    models_path="./",
):
    """Initializes a list of models

    Initializes a list of models where each model is the subdirectory
    beneath the given model root directory. The global list of models
    will exist on each processor.

    Parameters
    ----------
    model_root : str
        the directory whose subdirectories will become the model results
    models : list of str, optional
        only initialize a model whose name is in this list
    verbose : bool, optional
        enable to print information to the screen
    model_year : 2-tuple, optional
        shift model years from the first to the second part of the tuple

    Returns
    -------
    M : list of ILAMB.ModelResults.ModelsResults
       a list of the model results, sorted alphabetically by name

    """
    # initialize the models
    M = []
    if len(model_year) != 2:
        model_year = None
    max_model_name_len = 0
    if rank == 0 and verbose:
        print("\nSearching for model results in %s\n" % model_root)
    for subdir, dirs, files in os.walk(model_root):
        for mname in dirs:
            if len(models) > 0 and mname not in models:
                continue
            pkl_file = os.path.join(models_path, "%s.pkl" % mname)
            if os.path.isfile(pkl_file):
                with open(pkl_file, "rb") as infile:
                    m = pickle.load(infile)
            else:
                try:
                    m = ModelResult(
                        os.path.join(subdir, mname),
                        modelname=mname,
                        filter=filter,
                        regex=regex,
                        model_year=model_year,
                    )
                except Exception as ex:
                    if log:
                        logger.debug("[%s]" % mname, format_exc())
                    continue
            M.append(m)
            max_model_name_len = max(max_model_name_len, len(mname))
        break
    M = sorted(M, key=lambda m: m.name.upper())

    # assign unique colors
    clrs = il.GenerateDistinctColors(len(M))
    for m in M:
        m.color = clrs.pop(0)

    # save model objects as pickle files
    comm.Barrier()
    if rank == 0:
        for m in M:
            pkl_file = os.path.join(models_path, "%s.pkl" % m.name)
            with open(pkl_file, "wb") as out:
                pickle.dump(m, out, pickle.HIGHEST_PROTOCOL)

    # optionally output models which were found
    if rank == 0 and verbose:
        for m in M:
            print(("    {0:>45}").format(m.name))

    if len(M) == 0:
        if verbose and rank == 0:
            print("No model results found")
        comm.Barrier()
        comm.Abort(0)

    return M


def _parse_model_yaml(filename: str, cache_path: str = "./", only_models: list = []):
    """Setup models using a yaml file."""
    model_classes = {
        "ModelPointResult": ModelPointResult,
        "E3SMResult": E3SMResult,
        "ModelResult": ModelResult,
    }
    models = []
    with open(filename, encoding="utf-8") as fin:
        yml = yaml.safe_load(fin)
    for name, opts in yml.items():
        # optionally filter models
        if len(only_models) > 0 and name not in only_models:
            continue

        if "name" not in opts:
            opts["name"] = name

        # if the model_year option is given, convert to lits of floats
        if "model_year" in opts:
            opts["model_year"] = [
                float(y.strip()) for y in opts["model_year"].split(",")
            ]

        # select the class type
        cls = model_classes[opts["type"]] if "type" in opts else ModelResult
        if cls is None:
            typ = opts["type"]
            raise ValueError(f"The model type '{typ}' is not available")
        fcns = dir(cls)

        # if the pickle file exists, just load it
        cache = os.path.join(cache_path, f"{name}.pkl")
        if os.path.exists(cache):
            if "read_pickle" in fcns:
                model = cls().read_pickle(cache)
            else:
                with open(cache, mode="rb") as fin:
                    model = pickle.load(fin)
            models.append(model)
            continue

        # call the constructor using keywords defined in the YAML file
        cls = model_classes[opts["type"]] if "type" in opts else ModelResult
        model = cls(
            **{
                key: opts[key]
                for key in inspect.getfullargspec(cls).args
                if key in opts
            }
        )

        # some model types have a find_files() method, call if present loading
        # proper keywords from the YAML file
        if "find_files" in fcns:
            model.find_files(
                **{
                    key: opts[key]
                    for key in inspect.getfullargspec(model.find_files).args
                    if key in opts
                }
            )

        # some model types allow you to specify snynonms
        if "add_synonym" in fcns and "synonyms" in opts:
            for mvar, syn in opts["synonyms"].items():
                model.add_synonym(mvar, syn)

        # cache the model result
        if rank == 0:
            if "read_pickle" in fcns:
                model.to_pickle(cache)
            else:
                with open(cache, mode="wb") as fin:
                    pickle.dump(model, fin)

        models.append(model)

    for model in models:
        if isinstance(model.color, str) and model.color.startswith("#"):
            model.color = clr.hex2color(model.color)
    return models


def ParseModelSetup(
    model_setup, models=[], verbose=False, filter="", regex="", models_path="./"
):
    """Initializes a list of models

    Initializes a list of models where each model is the subdirectory
    beneath the given model root directory. The global list of models
    will exist on each processor.

    Parameters
    ----------
    model_setup : str
        the directory whose subdirectories will become the model results
    models : list of str, optional
        only initialize a model whose name is in this list
    verbose : bool, optional
        enable to print information to the screen

    Returns
    -------
    M : list of ILAMB.ModelResults.ModelsResults
       a list of the model results, sorted alphabetically by name

    """
    if rank == 0 and verbose:
        print("\nSetting up model results from %s\n" % model_setup)

    # intercept if this is a yaml file
    if model_setup.endswith(".yaml"):
        M = _parse_model_yaml(model_setup, cache_path=models_path, only_models=models)
        if rank == 0 and verbose:
            for m in M:
                print(("    {0:>45}").format(m.name))
            if len(M) == 0:
                print("No model results found")
                comm.Barrier()
                comm.Abort(0)
        return M

    # initialize the models
    M = []
    max_model_name_len = 0
    with open(model_setup) as f:
        for line in f.readlines():
            if line.strip().startswith("#"):
                continue
            line = line.split(",")
            mname = None
            mdir = None
            model_year = None
            mgrp = ""
            if len(line) >= 2:
                mname = line[0].strip()
                mdir = line[1].strip()
                # if mdir not a directory, then maybe path is relative to ILAMB_ROOT
                if not os.path.isdir(mdir):
                    mdir = os.path.join(os.environ["ILAMB_ROOT"], mdir).strip()
                if len(line) == 3:
                    mgrp = line[2].strip()
            if len(line) == 4:
                model_year = [float(line[2].strip()), float(line[3].strip())]
            max_model_name_len = max(max_model_name_len, len(mname))
            if (len(models) > 0 and mname not in models) or (mname is None):
                continue
            pkl_file = os.path.join(models_path, "%s.pkl" % mname)
            if os.path.isfile(pkl_file):
                with open(pkl_file, "rb") as infile:
                    m = pickle.load(infile)
            else:
                try:
                    m = ModelResult(
                        mdir,
                        modelname=mname,
                        filter=filter,
                        regex=regex,
                        model_year=model_year,
                        group=mgrp,
                    )
                except Exception as ex:
                    logger.debug("[%s]" % mname, format_exc())
                    continue
            M.append(m)

    # assign unique colors
    clrs = il.GenerateDistinctColors(len(M))
    for m in M:
        m.color = clrs.pop(0)

    # save model objects as pickle files
    comm.Barrier()
    if rank == 0:
        for m in M:
            pkl_file = os.path.join(models_path, "%s.pkl" % m.name)
            with open(pkl_file, "wb") as out:
                pickle.dump(m, out, pickle.HIGHEST_PROTOCOL)

    # optionally output models which were found
    if rank == 0 and verbose:
        for m in M:
            print(("    {0:>45}").format(m.name))

    if len(M) == 0:
        if verbose and rank == 0:
            print("No model results found")
        comm.Barrier()
        comm.Abort(0)

    return M


def InitializeRegions(filenames):
    """Initialize regions from a list of files.

    If the file is a netCDF4 file, see documentation in
    ILAMB.Regions.addRegionNetCDF4 for details on the required
    format. If the file defines regions by latitude/longitude bounds,
    then we anticipate comma delimited rows in the following form:

    shortname, longname, min lat, max lat, min lon, max lon

    Note that latitudes should be on (-90,90) and longitudes on
    (-180,180).

    Parameters
    ----------
    filenames : list of str
        a list of files from which to search for regions

    """
    r = Regions()
    for filename in filenames:
        try:
            r.addRegionNetCDF4(filename)
        except IOError:
            for line in open(filename):
                line = line.strip()
                if line.startswith("#"):
                    continue
                line = line.split(",")
                if len(line) == 6:
                    r.addRegionLatLonBounds(
                        line[0].strip(),
                        line[1].strip(),
                        [float(line[2]), float(line[3])],
                        [float(line[4]), float(line[5])],
                    )


def MatchRelationshipConfrontation(C):
    """Match relationship strings to confrontation longnames

    We allow for relationships to be studied by specifying the
    confrontation longname in the configure file. This routine loops
    over all defined relationships and finds the matching
    confrontation. (NOTE: this really belongs inside the Scoreboard
    object)

    Parameters
    ----------
    C : list of ILAMB.Confrontation.Confrontation
        the confrontation list

    Returns
    -------
    C : list of ILAMB.Confrontation.Confrontation
        the same list with relationships linked to confrontations
    """
    for c in C:
        if c.relationships is None:
            continue
        for i, longname in enumerate(c.relationships):
            found = False
            for cor in C:
                if longname.lower() == cor.longname.lower():
                    c.relationships[i] = cor
                    found = True
    return C


def FilterConfrontationList(C, match_list):
    """Filter the confrontation list

    Filter the confrontation list by requiring that at least one
    string in the input list is found in the longname in the
    confrontation.

    Parameters
    ----------
    C : list of ILAMB.Confrontation.Confrontation
       the source list of confrontations
    match_list : list of str
       the list of strings

    Returns
    -------
    Cf : list of ILAMB.Confrontation.Confrontation
        the list of filtered confrontations
    """
    if len(match_list) == 0:
        return C
    Cf = []
    for c in C:
        for match in match_list:
            if match in c.longname:
                Cf.append(c)
    return Cf


def BuildLocalWorkList(M, C, skip_cache=False):
    """Build the local work list

    We enumerate a list of work by taking combinations of model
    results and confrontations. This list is partitioned evenly among
    processes preferring to cluster as many confrontations with the
    same name together. While the work of the model-confrontation pair
    is local, some post-processing operations need performed once per
    confrontation. Thus we also need to flag one instance of each
    confrontation as the master process.

    Parameters
    ----------
    M : list of ILAMB.ModelResult.ModelResult
       list of models to analyze
    C : list of ILAMB.Confrontation.Confrontation
       list of confrontations

    Returns
    -------
    localW : list of (ILAMB.ModelResult.ModelResult, ILAMB.Confrontation.Confrontation) tuples
        the work local to this process
    """

    # Evenly divide up the work among processes
    W = []
    for c in C:
        for m in M:
            if skip_cache:
                # if we want to skip we have to check that it is complete
                fname = os.path.join(c.output_path, "%s_%s.nc" % (c.name, m.name))
                complete = False
                if os.path.isfile(fname):
                    try:
                        with Dataset(fname) as dset:
                            if "complete" in dset.ncattrs():
                                if dset.complete:
                                    complete = True
                    except:
                        pass
                if not complete:
                    os.system("rm -f %s" % fname)
                    W.append([m, c])
            else:
                W.append([m, c])

    wpp = float(len(W)) / size
    begin = int(round(rank * wpp))
    end = int(round((rank + 1) * wpp))
    localW = W[begin:end]

    # Determine who is the master of each confrontation
    for c in C:
        sendbuf = np.zeros(size, dtype="int")
        for w in localW:
            if c is w[1]:
                sendbuf[rank] += 1
        recvbuf = None
        if rank == 0:
            recvbuf = np.empty([size, sendbuf.size], dtype="int")
        comm.Gather(sendbuf, recvbuf, root=0)
        if rank == 0:
            numc = recvbuf.sum(axis=1)
        else:
            numc = np.empty(size, dtype="int")
        comm.Bcast(numc, root=0)
        if rank == numc.argmax():
            c.master = True
        else:
            c.master = False

    return localW


def WorkConfront(W, verbose=False, clean=False):
    """Performs the confrontation analysis

    For each model-confrontation pair (m,c) in the input work list,
    this routine will call c.confront(m) and keep track of the time
    required as well as any exceptions which are thrown.

    Parameters
    ----------
    W : list of (ILAMB.ModelResult.ModelResult, ILAMB.Confrontation.Confrontation) tuples
        the list of work
    verbose : bool, optional
        enable to print output to the screen monitoring progress
    clean : bool, optional
        enable to perform the confrontation again, overwriting previous results

    """
    maxCL = 45
    maxML = 20

    # Run analysis on your local work model-confrontation pairs
    for i, w in enumerate(W):
        m, c = w

        # if the results file exists, skip this confrontation unless we want to clean
        if (
            os.path.isfile(os.path.join(c.output_path, "%s_%s.nc" % (c.name, m.name)))
            and clean is False
        ):
            if verbose:
                print(
                    (
                        "    {0:>%d} {1:<%d} %sUsingCachedData%s "
                        % (maxCL, maxML, OK, ENDC)
                    ).format(c.longname, m.name)
                )
                sys.stdout.flush()
            continue

        # try to run the confrontation
        try:
            t0 = time.time()
            c.confront(m)
            dt = time.time() - t0
            proc[rank] += dt
            if verbose:
                dt = datetime.timedelta(seconds=max(1, int(np.round(dt))))
                print(
                    (
                        "    {0:>%d} {1:<%d} %sCompleted%s {2:>8}"
                        % (maxCL, maxML, OK, ENDC)
                    ).format(c.longname, m.name, str(dt))
                )
                sys.stdout.flush()

        # if things do not work out, print the exception so the user has some idea
        except Exception as ex:
            ierr[rank] = 1
            logger.debug("[%s][%s]\n%s" % (c.longname, m.name, format_exc()))
            if verbose:
                print(
                    (
                        "    {0:>%d} {1:<%d} %s%s%s"
                        % (maxCL, maxML, FAIL, ex.__class__.__name__, ENDC)
                    ).format(c.longname, m.name)
                )


def WorkPost(M, C, W, S, verbose=False, skip_plots=False):
    """Performs the post-processing

    Determines plot limits across all models, makes plots, generates
    other forms of HTML output.

    Parameters
    ----------
    M : list of ILAMB.ModelResult.ModelResult
       list of models to analyze
    C : list of ILAMB.Confrontation.Confrontation
       list of confrontations
    W : list of (ILAMB.ModelResult.ModelResult, ILAMB.Confrontation.Confrontation) tuples
        the list of work
    S : ILAMB.Scoreboard.Scoreboard
        the scoreboard context
    verbose : bool, optional
        enable to print output to the screen monitoring progress
    skip_plots : bool, optional
        enable to skip plotting
    """
    maxCL = 45
    maxML = 20
    for c in C:
        c.determinePlotLimits()
    for i, w in enumerate(W):
        m, c = w
        try:
            t0 = time.time()
            c.modelPlots(m)
            c.sitePlots(m)
            c.computeOverallScore(m)
            dt = time.time() - t0
            proc[rank] += dt
            if verbose:
                dt = datetime.timedelta(seconds=max(1, int(np.round(dt))))
                print(
                    (
                        "    {0:>%d} {1:<%d} %sCompleted%s {2:>8}"
                        % (maxCL, maxML, OK, ENDC)
                    ).format(c.longname, m.name, str(dt))
                )
                sys.stdout.flush()
        except Exception as ex:
            ierr[rank] = 1
            logger.debug("[%s][%s]\n%s" % (c.longname, m.name, format_exc()))
            if verbose:
                print(
                    (
                        "    {0:>%d} {1:<%d} %s%s%s"
                        % (maxCL, maxML, FAIL, ex.__class__.__name__, ENDC)
                    ).format(c.longname, m.name)
                )
                sys.stdout.flush()

    sys.stdout.flush()
    comm.Barrier()

    for i, c in enumerate(C):
        try:
            c.compositePlots()
        except Exception as ex:
            ierr[rank] = 1
            logger.debug("[compositePlots][%s]\n%s" % (c.longname, format_exc()))
        c.generateHtml()

    sys.stdout.flush()
    comm.Barrier()


def RestrictiveModelExtents(M, eps=2.0):
    extents0 = np.asarray([[-90.0, +90.0], [-180.0, +180.0]])
    extents = extents0.copy()
    for m in M:
        if not hasattr(m, "extents"):
            continue
        for i in range(2):
            extents[i, 0] = max(extents[i, 0], m.extents[i, 0])
            extents[i, 1] = min(extents[i, 1], m.extents[i, 1])
    diff = np.abs(extents0 - extents)
    extents = (diff <= eps) * extents0 + (diff > eps) * extents
    return extents


class MPIFileHandler(logging.FileHandler):
    """
    Class written by Di Cheng for parallel logging.

    https://gist.github.com/chengdi123000/42ec8ed2cbef09ee050766c2f25498cb

    """

    def __init__(
        self,
        filename,
        mode=MPI.MODE_WRONLY | MPI.MODE_CREATE | MPI.MODE_APPEND,
        encoding="utf-8",
        delay=False,
        comm=MPI.COMM_WORLD,
    ):
        self.baseFilename = os.path.abspath(filename)
        self.mode = mode
        self.encoding = encoding
        self.comm = comm
        if delay:
            # We don't open the stream, but we still need to call the
            # Handler constructor to set level, formatter, lock etc.
            logging.Handler.__init__(self)
            self.stream = None
        else:
            logging.StreamHandler.__init__(self, self._open())

    def _open(self):
        stream = MPI.File.Open(self.comm, self.baseFilename, self.mode)
        stream.Set_atomicity(True)
        return stream

    def emit(self, record):
        """
        Emit a record.
        If a formatter is specified, it is used to format the record.
        The record is then written to the stream with a trailing newline.  If
        exception information is present, it is formatted using
        traceback.print_exception and appended to the stream.  If the stream
        has an 'encoding' attribute, it is used to determine how to do the
        output to the stream.

        Modification:
            stream is MPI.File, so it must use `Write_shared` method rather
            than `write` method. And `Write_shared` method only accept
            bytestring, so `encode` is used. `Write_shared` should be invoked
            only once in each all of this emit function to keep atomicity.
        """
        try:
            msg = self.format(record)
            stream = self.stream
            stream.Write_shared((msg + self.terminator).encode(self.encoding))
            # self.flush()
        except Exception:
            self.handleError(record)

    def close(self):
        if self.stream:
            self.stream.Sync()
            self.stream.Close()
            self.stream = None


def ParseRunOptions(filename):
    run_opts = {}
    for line in open(filename).readlines():
        line = line.strip()
        if line.startswith("#!"):
            m3 = re.search(r"#!(.*)=(.*)", line)
            if m3:
                keyword = m3.group(1).strip()
                value = m3.group(2).strip().replace('"', "")
                run_opts[keyword] = value
    return run_opts


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
    "--model_root",
    dest="model_root",
    metavar="root",
    type=str,
    nargs=1,
    default=["./"],
    help="root at which to search for models",
)
parser.add_argument(
    "--config",
    dest="config",
    metavar="config",
    type=str,
    nargs=1,
    help="path to configuration file to use",
)
parser.add_argument(
    "--models",
    dest="models",
    metavar="m",
    type=str,
    nargs="+",
    default=[],
    help="specify which models to run, list model names with no quotes and only separated by a space.",
)
parser.add_argument(
    "--model_year",
    dest="model_year",
    metavar="y0 yf",
    type=int,
    nargs="+",
    default=[],
    help='set to shift model years, "--model_year y0 yf" will shift years from y0 to yf',
)
parser.add_argument(
    "--study_limits",
    dest="study_limits",
    metavar="y0 yf",
    type=int,
    nargs="+",
    default=[],
    help='set study period, "--study_limits y0 yf" will limit run from y0 thru yf',
)
parser.add_argument(
    "--confrontations",
    dest="confront",
    metavar="c",
    type=str,
    nargs="+",
    default=[],
    help="specify which confrontations to run, list confrontation names with no quotes and only separated by a space.",
)
parser.add_argument(
    "--regions",
    dest="regions",
    metavar="r",
    type=str,
    nargs="+",
    default=["global"],
    help="specify which regions to compute over",
)
parser.add_argument(
    "--clean",
    dest="clean",
    action="store_true",
    help="enable to remove analysis files and recompute",
)
parser.add_argument(
    "--disable_logging", dest="logging", action="store_false", help="disables logging"
)
parser.add_argument(
    "-q",
    "--quiet",
    dest="quiet",
    action="store_true",
    help="enable to silence screen output",
)
parser.add_argument(
    "--filter",
    dest="filter",
    metavar="filter",
    type=str,
    nargs=1,
    default=[""],
    help="a string which much be in the model filenames",
)
parser.add_argument(
    "--regex",
    dest="regex",
    metavar="regex",
    type=str,
    nargs=1,
    default=[""],
    help="a regular expression which filenames must conform to in order to be included",
)
parser.add_argument(
    "--build_dir",
    dest="build_dir",
    metavar="build_dir",
    type=str,
    nargs=1,
    default=["./_build"],
    help="path of where to save the output",
)
parser.add_argument(
    "--define_regions",
    dest="define_regions",
    type=str,
    nargs="+",
    default=[],
    help="list files containing user-defined regions",
)
parser.add_argument(
    "--model_setup",
    dest="model_setup",
    type=str,
    nargs="+",
    default=None,
    help="list files model setup information",
)
parser.add_argument(
    "--skip_plots",
    dest="skip_plots",
    action="store_true",
    help="enable to skip the plotting phase",
)
parser.add_argument(
    "--rel_only",
    dest="rel_only",
    action="store_true",
    help="enable only display relative differences in overall scores",
)
parser.add_argument(
    "--mem_per_pair",
    dest="mem_per_pair",
    metavar="MEM",
    type=float,
    default=100000.0,
    help="maximum memory for IOMB model-confrontation pairs",
)
parser.add_argument(
    "--title",
    dest="run_title",
    metavar="title",
    type=str,
    nargs=1,
    help="title of the study to use in the HTML output",
)
parser.add_argument(
    "--rmse_score_basis",
    dest="rmse_score_basis",
    metavar="basis",
    type=str,
    default="cycle",
    help='base the RMSE score on the full time series with "series" or just the annual cycle with "cycle"',
)
parser.add_argument(
    "--df_errs",
    dest="df_errs",
    metavar="df_errs",
    type=str,
    default=None,
    help="the pandas dataframe with the quantiles to use in scoring.",
)
parser.add_argument(
    "-g",
    "--global_region",
    dest="global_region",
    type=str,
    default=None,
    help="the ILAMB region to be the default (global) analysis region.",
)
args = parser.parse_args()
if args.config is None:
    if rank == 0:
        print(
            "\nError: You must specify a configuration file using the option --config\n"
        )
    comm.Barrier()
    comm.Abort(1)

# Additional options could be in the configure file
run_opts = ParseRunOptions(args.config[0])
for key in run_opts:
    if key in ["define_regions"]:
        define_regions = [
            os.path.join(os.environ["ILAMB_ROOT"], r) for r in run_opts[key].split(",")
        ]
        args.__dict__[key] += define_regions
assert args.rmse_score_basis in ["series", "cycle"]

# Setup regions
r = Regions()
InitializeRegions(args.define_regions)
missing = []
for region in args.regions:
    if region not in r.regions:
        missing.append(region)
if len(missing) > 0:
    raise ValueError(
        "Unable to find the following regions %s from the following list of possible regions %s"
        % (missing, r.regions)
    )
if args.global_region:
    r.setGlobalRegion(args.global_region)

# Setup study
T0 = time.time()
if rank == 0:
    if not os.path.isdir(args.build_dir[0]):
        os.makedirs(args.build_dir[0])
if args.model_setup is None:
    M = InitializeModels(
        args.model_root[0],
        args.models,
        not args.quiet,
        filter=args.filter[0],
        regex=args.regex[0],
        model_year=args.model_year,
        models_path=args.build_dir[0],
    )
else:
    M = ParseModelSetup(
        args.model_setup[0],
        args.models,
        not args.quiet,
        filter=args.filter[0],
        models_path=args.build_dir[0],
    )


try:
    df_errs = None
    if args.df_errs is not None:
        df_errs = pd.read_parquet(args.df_errs)
except:
    if rank == 0:
        print("Unable to read quantiles")
    comm.Abort(0)


if rank == 0 and not args.quiet:
    print("\nParsing config file %s...\n" % args.config[0])
S = Scoreboard(
    args.config[0],
    regions=args.regions,
    master=rank == 0,
    verbose=not args.quiet,
    build_dir=args.build_dir[0],
    extents=RestrictiveModelExtents(M),
    rel_only=args.rel_only,
    mem_per_pair=args.mem_per_pair,
    run_title=args.run_title,
    rmse_score_basis=args.rmse_score_basis,
    df_errs=df_errs,
)
C = MatchRelationshipConfrontation(S.list())
if len(args.study_limits) == 2:
    args.study_limits[1] += 1
    for c in C:
        c.study_limits = (np.asarray(args.study_limits) - 1850) * 365.0
Cf = FilterConfrontationList(C, args.confront)
if rank == 0:
    os.system(
        "cp %s %s" % (args.config[0], os.path.join(args.build_dir[0], "ilamb.cfg"))
    )

# Setup logging
logger = logging.getLogger("%i" % comm.rank)
logname = ""
formatter = logging.Formatter("[%(levelname)s][%(name)s][%(funcName)s]%(message)s")
logger.setLevel(logging.DEBUG)
if args.logging:
    logname = "%s/ILAMB%02d.log" % (
        S.build_dir,
        len(glob.glob("%s/*.log" % S.build_dir)) + 1,
    )
    mh = MPIFileHandler(logname)
    mh.setFormatter(formatter)
    logger.addHandler(mh)

if rank == 0:
    logger.info(" " + " ".join(platform.uname()))
    for key in [
        "ILAMB",
        "numpy",
        "matplotlib",
        "netCDF4",
        "cf_units",
        "sympy",
        "mpi4py",
    ]:
        pkg = __import__(key)
        try:
            path = pkg.__path__[0]
        except:
            path = key
        logger.info(" %s (%s)" % (path, pkg.__version__))
    logger.info(" %s" % datetime.datetime.now())

if rank == 0 and not args.quiet and len(Cf) != len(C):
    print("\nWe filtered some confrontations, actually running...\n")
    for c in Cf:
        print(("    {0:>45}").format(c.longname))
C = Cf

sys.stdout.flush()
comm.Barrier()

if rank == 0 and not args.quiet:
    print("\nRunning model-confrontation pairs...\n")

sys.stdout.flush()
comm.Barrier()

W = BuildLocalWorkList(M, C, skip_cache=True)
WorkConfront(W, not args.quiet, args.clean)

sys.stdout.flush()
comm.Barrier()

if not args.skip_plots:
    if rank == 0 and not args.quiet:
        print("\nFinishing post-processing which requires collectives...\n")

    sys.stdout.flush()
    comm.Barrier()

    W = BuildLocalWorkList(M, C, skip_cache=False)
    WorkPost(M, C, W, S, not args.quiet)

if rank == 0:
    S.createHtml(M)
    S.createUDDashboard()

sys.stdout.flush()
comm.Barrier()

# Runtime information
proc_reduced = np.zeros(proc.shape)
ierr_reduced = np.zeros(ierr.shape)
comm.Reduce(proc, proc_reduced, root=0)
comm.Reduce(ierr, ierr_reduced, root=0)
if size > 1:
    logger.info("[process time] %.1f s" % proc[rank])
if rank == 0:
    logger.info("[total time] %.1f s" % (time.time() - T0))
    if size > 1:
        if proc_reduced.min() > 1e-6:
            logger.info(
                "[process balance] %.2f" % (proc_reduced.max() / proc_reduced.min())
            )
        else:
            logger.info("[process balance] nan")
        logger.info(
            "[parallel efficiency] %.0f%%"
            % (100.0 * proc_reduced.sum() / float(size) / (time.time() - T0))
        )

if rank == 0:
    S.dumpScores(M, "scores.csv")
    S.harvestInformation(M)

if rank == 0 and ierr_reduced.max() > 0 and args.logging and not args.quiet:
    print(
        "\nErrors occurred in the run, please consult %s for more detailed information"
        % logname
    )

if rank == 0 and not args.quiet:
    print(
        "\nCompleted in {0:>8}\n".format(
            str(datetime.timedelta(seconds=int(np.round((time.time() - T0)))))
        )
    )
