# -*- coding: utf-8 -*-
"""Coastal Blue Carbon Preprocessor."""
import logging
import os
import time

import pygeoprocessing
import taskgraph
from osgeo import gdal

from .. import gettext
from .. import spec_utils
from .. import utils
from .. import validation
from ..model_metadata import MODEL_METADATA
from ..unit_registry import u
from . import coastal_blue_carbon

LOGGER = logging.getLogger(__name__)

BIOPHYSICAL_COLUMNS_SPEC = coastal_blue_carbon.MODEL_SPEC[
    'args']['biophysical_table_path']['columns']

MODEL_SPEC = {
    "model_id": "coastal_blue_carbon_preprocessor",
    "model_name": MODEL_METADATA["coastal_blue_carbon_preprocessor"].model_title,
    "pyname": MODEL_METADATA["coastal_blue_carbon_preprocessor"].pyname,
    "userguide": MODEL_METADATA["coastal_blue_carbon_preprocessor"].userguide,
    "args": {
        "workspace_dir": spec_utils.WORKSPACE,
        "results_suffix": spec_utils.SUFFIX,
        "n_workers": spec_utils.N_WORKERS,
        "lulc_lookup_table_path": {
            "name": gettext("LULC lookup table"),
            "type": "csv",
            "about": gettext(
                "A table mapping LULC codes from the snapshot rasters to the "
                "corresponding LULC class names, and whether or not the "
                "class is a coastal blue carbon habitat."),
            "index_col": "lucode",
            "columns": {
                "lucode": {
                    "type": "integer",
                    "about": gettext(
                        "LULC code. Every value in the "
                        "snapshot LULC maps must have a corresponding entry "
                        "in this column.")},
                "lulc-class": {
                    "type": "freestyle_string",
                    "about": gettext("Name of the LULC class.")},
                "is_coastal_blue_carbon_habitat": {
                    "type": "boolean",
                    "about": gettext(
                        "Enter TRUE if this LULC class is a coastal blue "
                        "carbon habitat, FALSE if not.")}
            }
        },
        "landcover_snapshot_csv": {
            "type": "csv",
            "index_col": "snapshot_year",
            "columns": {
                "snapshot_year": {
                    "type": "integer",
                    "about": gettext("Year to snapshot.")},
                "raster_path": {
                    "type": "raster",
                    "bands": {1: {"type": "integer"}},
                    "about": gettext(
                        "Map of LULC in the snapshot year. "
                        "All values in this raster must have corresponding "
                        "entries in the LULC Lookup table.")
                }
            },
            "about": gettext(
                "A table mapping snapshot years to corresponding LULC maps "
                "for each year."),
            "name": gettext("LULC snapshots table")
        }
    },
    "outputs": {
        "transitions.csv": {
            "about": (
                "LULC transition matrix. The first column represents the "
                "source LULC class, and the first row represents the "
                "destination LULC classes. Cells are populated with "
                "transition states, or left empty if no such transition occurs."),
            "index_col": "lulc-class",
            "columns": {
                "lulc-class": {
                    "type": "freestyle_string",
                    "about": gettext(
                        "LULC class names matching those in the biophysical "
                        "table.")},
                "[LULC]": {
                    "type": "option_string",
                    "options": {
                        "accum": {
                            "description": gettext("a state of carbon accumulation")
                        },
                        "disturb": {
                            "description": gettext(
                                "Carbon disturbance occurred. Replace this "
                                "with one of ‘low-impact-disturb’, "
                                "‘med-impact-disturb’, or ‘high-impact-disturb’ "
                                "to indicate the degree of disturbance.")},
                        "NCC": {
                            "description": gettext("no change in carbon")
                        }
                    }
                }
            }
        },
        "carbon_pool_transient_template.csv": {
            "about": (
                "Table mapping each LULC type to impact and accumulation "
                "information. This is a template that you will fill out to "
                "create the biophysical table input to the main model."),
            "index_col": "lucode",
            "columns": {
                **BIOPHYSICAL_COLUMNS_SPEC,
                # remove "expression" property which doesn't go in output spec
                "biomass-half-life": dict(
                    set(BIOPHYSICAL_COLUMNS_SPEC["biomass-half-life"].items()) -
                    {("expression", "value > 0")}
                ),
                "soil-half-life": dict(
                    set(BIOPHYSICAL_COLUMNS_SPEC["soil-half-life"].items()) -
                    {("expression", "value > 0")}
                )
            }
        },
        "aligned_lulc_[YEAR].tif": {
            "about": (
                "Copy of LULC map for the given year, aligned and resampled "
                "to match all the other LULC maps."),
            "bands": {1: {"type": "integer"}}
        },
        "taskgraph_cache": spec_utils.TASKGRAPH_DIR
    }
}


ALIGNED_LULC_RASTER_TEMPLATE = 'aligned_lulc_{year}{suffix}.tif'
TRANSITION_TABLE = 'carbon_pool_transition_template{suffix}.csv'
BIOPHYSICAL_TABLE = 'carbon_biophysical_table_template{suffix}.csv'


def execute(args):
    """Coastal Blue Carbon Preprocessor.

    The preprocessor accepts a list of rasters and checks for cell-transitions
    across the rasters.  The preprocessor outputs a CSV file representing a
    matrix of land cover transitions, each cell pre-filled with a string
    indicating whether carbon accumulates or is disturbed as a result of the
    transition, if a transition occurs.

    Args:
        args['workspace_dir'] (string): directory path to workspace
        args['results_suffix'] (string): append to outputs directory name if
            provided
        args['lulc_lookup_table_path'] (string): filepath of lulc lookup table
        args['landcover_csv_path'] (string): filepath to a CSV containing the
            year and filepath to snapshot rasters on disk.  The years may be in
            any order, but must be unique.

    Returns:
        ``None``
    """
    suffix = utils.make_suffix_string(args, 'results_suffix')
    output_dir = os.path.join(args['workspace_dir'], 'outputs_preprocessor')
    utils.make_directories([output_dir])

    try:
        n_workers = int(args['n_workers'])
    except (KeyError, ValueError, TypeError):
        # KeyError when n_workers is not present in args
        # ValueError when n_workers is an empty string.
        # TypeError when n_workers is None.
        n_workers = -1  # Synchronous mode.
    task_graph = taskgraph.TaskGraph(
        os.path.join(args['workspace_dir'], 'taskgraph_cache'),
        n_workers, reporting_interval=5.0)

    snapshots_dict = validation.get_validated_dataframe(
        args['landcover_snapshot_csv'],
        **MODEL_SPEC['args']['landcover_snapshot_csv']
    )['raster_path'].to_dict()

    # Align the raster stack for analyzing the various transitions.
    min_pixel_size = float('inf')
    source_snapshot_paths = []
    aligned_snapshot_paths = []
    for snapshot_year, raster_path in sorted(
            snapshots_dict.items(), key=lambda x: x[0]):
        source_snapshot_paths.append(raster_path)
        aligned_snapshot_paths.append(os.path.join(
            output_dir, ALIGNED_LULC_RASTER_TEMPLATE.format(
                year=snapshot_year, suffix=suffix)))
        min_pixel_size = min(
            utils.mean_pixel_size_and_area(
                pygeoprocessing.get_raster_info(raster_path)['pixel_size'])[0],
            min_pixel_size)

    baseline_srs_wkt = pygeoprocessing.get_raster_info(
        snapshots_dict[min(snapshots_dict.keys())])['projection_wkt']
    alignment_task = task_graph.add_task(
        func=pygeoprocessing.align_and_resize_raster_stack,
        args=(source_snapshot_paths,
              aligned_snapshot_paths,
              (['near']*len(source_snapshot_paths)),
              (min_pixel_size, -min_pixel_size),
              'intersection'),
        kwargs={'target_projection_wkt': baseline_srs_wkt},
        target_path_list=aligned_snapshot_paths,
        task_name='Align input landcover rasters')

    landcover_df = validation.get_validated_dataframe(
        args['lulc_lookup_table_path'],
        **MODEL_SPEC['args']['lulc_lookup_table_path'])

    target_transition_table = os.path.join(
        output_dir, TRANSITION_TABLE.format(suffix=suffix))
    _ = task_graph.add_task(
        func=_create_transition_table,
        args=(landcover_df,
              aligned_snapshot_paths,
              target_transition_table),
        target_path_list=[target_transition_table],
        dependent_task_list=[alignment_task],
        task_name='Determine transitions and write transition table')

    target_biophysical_table_path = os.path.join(
        output_dir, BIOPHYSICAL_TABLE.format(suffix=suffix))
    _ = task_graph.add_task(
        func=_create_biophysical_table,
        args=(landcover_df, target_biophysical_table_path),
        target_path_list=[target_biophysical_table_path],
        task_name='Write biophysical table template')

    task_graph.close()
    task_graph.join()


def _create_transition_table(landcover_df, lulc_snapshot_list,
                             target_table_path):
    """Create the transition table from a series of landcover snapshots.

    Args:
        landcover_df (pandas.DataFrame: A table mapping integer landcover
            codes to values indicating the landcover class name in the
            ``lulc-class`` column and ``True`` or ``False`` under the
            ``is_coastal_blue_carbon_habitat`` column.
        lulc_snapshot_list (list): A list of string paths to GDAL rasters on
            disk.  All rasters must have the same spatial reference, pixel size
            and dimensions and must also all be integer rasters, where all
            non-nodata pixel values must be represented in the
            ``landcover_df`` dataframe.
        target_table_path (string): A string path to where the target
            transition table should be written.

    Returns:
        ``None``.
    """
    n_rows, n_cols = pygeoprocessing.get_raster_info(
        lulc_snapshot_list[0])['raster_size']
    n_pixels_total = (n_rows * n_cols) * len(lulc_snapshot_list)
    n_pixels_processed = 0

    raster_tuple_list = []
    for raster_path in lulc_snapshot_list:
        raster = gdal.OpenEx(raster_path, gdal.OF_RASTER)
        band = raster.GetRasterBand(1)
        nodata = band.GetNoDataValue()
        raster_tuple_list.append((raster, band, nodata))

    transition_pairs = set()
    last_log_time = time.time()
    for block_offsets in pygeoprocessing.iterblocks((lulc_snapshot_list[0], 1),
                                                    offset_only=True):
        _, from_band, from_nodata = raster_tuple_list[0]
        from_array = from_band.ReadAsArray(**block_offsets)
        from_band = None

        for (_, to_band, to_nodata) in raster_tuple_list[1:]:
            if time.time() - last_log_time >= 5.0:
                percent_complete = n_pixels_processed / n_pixels_total
                LOGGER.info(
                    "Determining landcover transitions, "
                    f"{percent_complete:.2f}% complete.")

            to_array = to_band.ReadAsArray(**block_offsets)

            # This comparison assumes that our landcover rasters are of an
            # integer type.  When int matrices, we can compare directly to
            # None.
            valid_pixels = (
                ~pygeoprocessing.array_equals_nodata(from_array, from_nodata) &
                ~pygeoprocessing.array_equals_nodata(to_array, to_nodata))
            transition_pairs = transition_pairs.union(
                set(zip(from_array[valid_pixels].flatten(),
                        to_array[valid_pixels].flatten())))

            # Swap the arrays around to use the current 'to_array', 'to_nodata'
            # as the 'from_array', 'from_nodata' in the next iteration.
            from_array, from_nodata = (to_array, to_nodata)
            n_pixels_processed += to_array.size
    raster_tuple_list = None
    to_band = None
    LOGGER.info("Determining landcover transitions, 100.00%% complete.")

    # Mapping of whether the from, to landcover types are coastal blue carbon
    # habitats to the string carbon transition type.
    # The keys are structured as a tuple of two booleans where:
    #  * tuple[0] = whether the FROM transition is CBC habitat
    #  * tuple[1] = whether the TO transition is CBC habitat
    transition_types = {
        (True, True): 'accum',  # veg --> veg
        (False, True): 'accum',  # non-veg --> veg
        (True, False): 'disturb',  # veg --> non-veg
        (False, False): 'NCC',  # non-veg --> non-veg
    }

    sparse_transition_table = {}
    for from_lucode, to_lucode in transition_pairs:
        try:
            from_is_cbc = landcover_df[
                'is_coastal_blue_carbon_habitat'][from_lucode]
            to_is_cbc = landcover_df[
            'is_coastal_blue_carbon_habitat'][to_lucode]
        except KeyError:
            for variable in (from_lucode, to_lucode):
                if variable not in landcover_df.index:
                    raise ValueError(
                        'The landcover table is missing a row with the '
                        f'landuse code {variable}.')

        sparse_transition_table[(from_lucode, to_lucode)] = (
            transition_types[(from_is_cbc, to_is_cbc)])

    code_list = sorted(landcover_df.index)
    lulc_class_list_sorted = [
        landcover_df['lulc-class'][code] for code in code_list]
    with open(target_table_path, 'w') as csv_file:
        fieldnames = ['lulc-class'] + lulc_class_list_sorted
        csv_file.write(f"{','.join(fieldnames)}\n")
        for row_code in code_list:
            class_name = landcover_df['lulc-class'][row_code]
            row = [class_name]
            for col_code in code_list:
                try:
                    column_value = sparse_transition_table[
                        (row_code, col_code)]
                except KeyError:
                    # When there isn't a transition that we know about, just
                    # leave the table blank.
                    column_value = ''
                row.append(column_value)
            csv_file.write(','.join(row) + '\n')

        # Append legend
        csv_file.write("\n,legend")
        csv_file.write(
            "\n,empty cells indicate that no transitions occur of that type")
        csv_file.write("\n,disturb (disturbance): change to low- med- or "
                       "high-impact-disturb")
        csv_file.write("\n,accum (accumulation)")
        csv_file.write("\n,NCC (no-carbon-change)")


def _create_biophysical_table(landcover_df, target_biophysical_table_path):
    """Write the biophysical table template to disk.

    The biophysical table templates contains all of the fields required by the
    main Coastal Blue Carbon model, and any field values that exist in the
    landcover table provided to this model will be carried over to the new
    table.

    Args:
        landcover_df (pandas.DataFrame): A table mapping int landcover codes
            to biophysical data
        target_biophysical_table_path (string): The path to where the
            biophysical table template will be stored on disk.

    Returns:
        ``None``
    """
    target_column_names = [
        colname.lower() for colname in coastal_blue_carbon.MODEL_SPEC['args'][
            'biophysical_table_path']['columns']]

    with open(target_biophysical_table_path, 'w') as bio_table:
        bio_table.write(f"{','.join(target_column_names)}\n")
        for lulc_code, row in landcover_df.sort_index().iterrows():
            # 2 columns are defined below, and we need 1 less comma to only
            # have commas between fields.
            row = []
            for colname in target_column_names:
                if colname == 'lucode':
                    row.append(str(lulc_code))
                else:
                    try:
                        # Use the user's defined value if it exists
                        row.append(str(landcover_df[colname][lulc_code]))
                    except KeyError:
                        row.append('')
            bio_table.write(f"{','.join(row)}\n")


@validation.invest_validator
def validate(args, limit_to=None):
    """Validate an input dictionary for Coastal Blue Carbon: Preprocessor.

    Args:
        args (dict): The args dictionary.
        limit_to=None (str or None): If a string key, only this args parameter
            will be validated.  If ``None``, all args parameters will be
            validated.

    Returns:
        A list of tuples where tuple[0] is an iterable of keys that the error
        message applies to and tuple[1] is the string validation warning.
    """
    return validation.validate(args, MODEL_SPEC['args'])
