"""Map with vulnerability classes."""

import os

from osgeo import gdal
import pandas as pd

from ..misc import progress_bar, makeblock


def vulnerability_map(
        forest_file,
        dist_file,
        dist_bins,
        subj_file,
        output_file="vulnerability_map.tif",
        blk_rows=128,
        verbose=True):
    """Map with vulnerability classes.

    A raster file with vulnerability classes is created combining distance
    to forest edge classes and subjurisdiction identifiers. High
    values indicate higher vulnerability. Raster type is UInt16 ([0,
    65535]). NoData value is set to 0.

    :param forest_file: Input forest cover file. Necessary to have
        forest extent within country borders (info not provided by
        dist_file).

    :param dist_file: Input file of distance to forest edge.

    :param dist_bins: The distance bins used to convert distance to
        forest edge into vulnerability classes. The first value
        indicates the minimal distance and the last value indicates
        the distance threshold.

    :param subj_file: Input raster file with subjurisdiction
        identifiers.

    :param blk_rows: If > 0, number of rows for computation by block.

    :param verbose: Logical. Whether to print messages or not. Default
        to ``True``.

    :return: Bins used to categorize the deforestation risk based on
        the distance to forest edge.

    """

    # ==============================================================
    # Input rasters
    # ==============================================================

    # Forest cover
    forest_ds = gdal.Open(forest_file)
    forest_band = forest_ds.GetRasterBand(1)
    # Raster size
    xsize = forest_band.XSize
    ysize = forest_band.YSize

    # Distance to forest edge raster file
    dist_ds = gdal.Open(dist_file)
    dist_band = dist_ds.GetRasterBand(1)

    # Subjurisdictions raster file
    subj_ds = gdal.Open(subj_file)
    subj_band = subj_ds.GetRasterBand(1)

    # Make blocks
    blockinfo = makeblock(forest_file, blk_rows=blk_rows)
    nblock = blockinfo[0]
    nblock_x = blockinfo[1]
    x = blockinfo[3]
    y = blockinfo[4]
    nx = blockinfo[5]
    ny = blockinfo[6]

    # Distance info from bins
    n_classes = len(dist_bins) - 1
    dist_thresh = dist_bins[-1]

    # =================================================
    # Categorical raster file for vulnerability classes
    # =================================================

    # Create categorical (cat) raster file for deforestation risk
    driver = gdal.GetDriverByName("GTiff")
    if os.path.isfile(output_file):
        os.remove(output_file)
    cat_ds = driver.Create(output_file, xsize, ysize, 1,
                           gdal.GDT_UInt16,
                           ["COMPRESS=DEFLATE", "PREDICTOR=2",
                            "BIGTIFF=YES"])
    cat_ds.SetProjection(dist_ds.GetProjection())
    cat_ds.SetGeoTransform(dist_ds.GetGeoTransform())
    cat_band = cat_ds.GetRasterBand(1)
    cat_band.SetNoDataValue(0)

    # =================
    # Categorizing
    # =================

    # Loop on blocks of data
    for b in range(nblock):
        # Progress bar
        if verbose:
            progress_bar(nblock, b + 1)
        # Position
        px = b % nblock_x
        py = b // nblock_x
        # Data
        forest_data = forest_band.ReadAsArray(x[px], y[py], nx[px], ny[py])
        dist_data = dist_band.ReadAsArray(x[px], y[py], nx[px], ny[py])
        subj_data = subj_band.ReadAsArray(x[px], y[py], nx[px], ny[py])
        # Categorize
        cat_data = pd.cut(dist_data.flatten(), bins=dist_bins,
                          labels=False, include_lowest=True,
                          right=True)
        cat_data = cat_data.reshape(dist_data.shape)
        cat_data = n_classes + 1 - cat_data
        # Set classes 1 beyond distance threshold
        cat_data[dist_data > dist_thresh] = 1
        cat_data[forest_data == 0] = 0  # Mask with forest
        # Adding subjurisdiction info
        cat_data = cat_data * 1000 + subj_data
        cat_data[cat_data <= 1000] = 0
        # Write to file
        cat_band.WriteArray(cat_data, x[px], y[py])

    # Compute statistics
    cat_band.FlushCache()
    cb = gdal.TermProgress_nocb if verbose else 0
    cat_band.ComputeStatistics(False, cb)

    # Dereference drivers
    forest_band = None
    cat_band = None
    dist_band = None
    subj_band = None
    del forest_ds, cat_ds, dist_ds, subj_ds


# # Test
# import os
# os.chdir("/home/ghislain/deforisk/MTQ_2000_2010_2020_jrc_7221/")
# bins = vulnerability_classes(
#     dist_file="data/dist_edge.tif",
#     dist_thresh=240,
#     subj_file="outputs/benchmark_model/subjurisdictions.tif",
#     output_file="outputs/benchmark_model/vulnerability_classes.tif",
#     blk_rows=128,
#     verbose=False)

# End
