"""
Calculate zonal statistics from vector and raster files.

BROKEN

TODO:
    - Fix rasterize step
    - Improve documentation

"""

import sys; sys.path.append("../../") # Path: buteo/vector/zonal_statistics.py
import numpy as np

from numba import jit
from osgeo import ogr

from buteo.raster.core_raster import raster_to_array, raster_to_metadata
from buteo.vector.rasterize import rasterize_vector
from buteo.vector.reproject import reproject_vector
from buteo.vector.core_vector import (
    open_vector,
    _vector_to_memory,
    _vector_to_metadata,
)
from buteo.vector.zonal_statistics_stats import calculate_array_stats
from buteo.utils.core_utils import progress



import numpy as np
from enum import Enum

from numba import jit


class stat(Enum):
    count = 1
    range = 2
    min = 3
    max = 4
    sum = 5
    mean = 6
    avg = 6
    average = 6
    var = 7
    variance = 7
    std = 8
    stdev = 8
    standard_deviation = 8
    skew = 9
    kurtosis = 10
    median = 11
    med = 11
    iqr = 12
    q02 = 13
    q98 = 14
    q1 = 15
    q3 = 16
    mad = 17
    median_absolute_deviation = 17
    mode = 18
    snr = 19
    eff = 20
    cv = 21


@jit(nopython=True, parallel=True, nogil=True, fastmath=True, inline="always")
def calculate_array_stats(arr, stats):
    stats_length = int(len(stats))
    result = np.zeros((stats_length), dtype="float32")

    for idx in range(stats_length):
        if stat(stats[idx]) == 1:  # count
            result[idx] = arr.size
        elif stat(stats[idx]) == 2:  # range
            result[idx] = np.ptp(arr)
        elif stat(stats[idx]) == 3:  # min
            result[idx] = np.min(arr)
        elif stat(stats[idx]) == 4:  # max
            result[idx] = np.max(arr)
        elif stat(stats[idx]) == 5:  # sum
            result[idx] = np.sum(arr)
        elif stat(stats[idx]) == 6:  # mean
            result[idx] = np.mean(arr)
        elif stat(stats[idx]) == 7:  # var
            result[idx] = np.var(arr)
        elif stat(stats[idx]) == 8:  # std
            result[idx] = np.std(arr)
        elif stat(stats[idx]) == 9:  # skew
            mean = np.mean(arr)
            std = np.std(arr)
            if std == 0:
                result[idx] = 0.0
                continue
            deviations = np.sum(np.power(arr - mean, 3))
            result[idx] = (deviations * (1 / arr.size)) / (np.power(std, 3))
        elif stat(stats[idx]) == 10:  # kurt
            mean = np.mean(arr)
            std = np.std(arr)
            if std == 0:
                result[idx] = 0.0
                continue
            deviations = np.sum(np.power(arr - mean, 4))
            result[idx] = (deviations * (1 / arr.size)) / (np.power(std, 4))
        elif stat(stats[idx]) == 11:  # median
            result[idx] = np.median(arr)
        elif stat(stats[idx]) == 12:  # iqr
            result[idx] = np.quantile(
                arr, np.array([0.25, 0.75], dtype="float32")
            ).sum()
        elif stat(stats[idx]) == 13:  # q02
            result[idx] = np.quantile(arr, 0.02)
        elif stat(stats[idx]) == 14:  # q92
            result[idx] = np.quantile(arr, 0.98)
        elif stat(stats[idx]) == 15:  # q1
            result[idx] = np.quantile(arr, 0.25)
        elif stat(stats[idx]) == 16:  # q3
            result[idx] = np.quantile(arr, 0.75)
        elif stat(stats[idx])== 17:  # mad
            median = np.median(arr)
            absdev = np.abs(arr - median)
            result[idx] = np.median(absdev)
        elif stat(stats[idx]) == 18:  # mode
            uniques = np.unique(arr)
            counts = np.zeros_like(uniques, dtype="uint64")

            # numba does not support return count of uniques.
            for idx_values in range(len(arr)):
                val = arr[idx_values]
                for idx_uniques in range(len(uniques)):
                    if val == uniques[idx_uniques]:
                        counts[idx_uniques] = counts[idx_uniques] + 1
            index = np.argmax(counts)
            result[idx] = uniques[index]
        elif stat(stats[idx]) == 19:  # snr
            std = np.std(arr)
            if std == 0:
                result[idx] = 0.0
                continue
            result[idx] = np.mean(arr) / std
        elif stat(stats[idx]) == 20:  # eff
            mean = np.mean(arr)
            if mean == 0:
                result[idx] = 0.0
                continue
            result[idx] = np.var(arr) / np.power(mean, 2)
        elif stat(stats[idx]) == 21:  # cv
            mean = np.mean(arr)
            if mean == 0:
                result[idx] = 0.0
                continue
            result[idx] = np.std(arr) / mean

    return result

@jit(nopython=True, nogil=True, fastmath=True, inline="always")
def overlap_size_calc(extent, raster_transform):
    return np.array(
        [
            (extent[1] - extent[0]) / raster_transform[1],
            (extent[3] - extent[2]) / abs(raster_transform[5]),
        ],
        dtype=np.int32,
    )


@jit(nopython=True, nogil=True, fastmath=True, inline="always")
def align_extent(raster_transform, vector_extent, raster_size):
    pixel_width = abs(raster_transform[1])
    pixel_height = abs(raster_transform[5])

    raster_min_x = raster_transform[0]
    # raster_max_x = raster_min_x + (raster_size[0] * pixel_width)
    raster_max_y = raster_transform[3]
    raster_min_y = raster_max_y + (raster_size[1] * -pixel_width)

    vector_min_x = vector_extent[0]
    vector_max_x = vector_extent[1]
    vector_min_y = vector_extent[2]
    vector_max_y = vector_extent[3]

    # Align the two extents
    vector_min_x = vector_min_x - (vector_min_x - raster_min_x) % pixel_width
    vector_max_x = vector_max_x + (vector_max_x - raster_min_x) % pixel_width
    vector_min_y = vector_min_y - (vector_min_y - raster_max_y) % pixel_height
    vector_max_y = vector_max_y + (vector_max_y - raster_max_y) % pixel_height

    rasterized_x_size = int((vector_max_x - vector_min_x) / pixel_width)
    rasterized_y_size = int((vector_max_y - vector_min_y) / pixel_height)

    rasterized_x_offset = int((vector_min_x - raster_min_x) / pixel_width)
    rasterized_y_offset = (
        int(raster_size[1] - int((vector_min_y - raster_min_y) / pixel_height))
        - rasterized_y_size
    )

    if rasterized_x_offset < 0:
        rasterized_x_offset = 0

    if rasterized_y_offset < 0:
        rasterized_y_offset = 0

    if (rasterized_x_offset + rasterized_x_size) > raster_size[0]:
        rasterized_x_offset = rasterized_x_offset - (
            (rasterized_x_offset + rasterized_x_size) - raster_size[0]
        )

    if (rasterized_y_offset + rasterized_y_size) > raster_size[1]:
        rasterized_y_offset = rasterized_y_offset - (
            (rasterized_y_offset + rasterized_y_size) - raster_size[1]
        )

    new_vector_extent = np.array(
        [vector_min_x, vector_max_x, vector_min_y, vector_max_y], dtype=np.float32
    )
    rasterized_size = np.array(
        [rasterized_x_size, rasterized_y_size, pixel_width, pixel_height],
        dtype=np.float32,
    )
    offset = np.array([rasterized_x_offset, rasterized_y_offset], dtype=np.int32)

    return new_vector_extent, rasterized_size, offset


@jit(nopython=True, nogil=True, fastmath=True, inline="always")
def get_intersection(extent1, extent2):
    one_bottomLeftX = extent1[0]
    one_topRightX = extent1[1]
    one_bottomLeftY = extent1[2]
    one_topRightY = extent1[3]

    two_bottomLeftX = extent2[0]
    two_topRightX = extent2[1]
    two_bottomLeftY = extent2[2]
    two_topRightY = extent2[3]

    if two_bottomLeftX > one_topRightX:  # Too far east
        return np.array([0, 0, 0, 0], dtype=np.float32)
    elif two_bottomLeftY > one_topRightY:  # Too far north
        return np.array([0, 0, 0, 0], dtype=np.float32)
    elif two_topRightX < one_bottomLeftX:  # Too far west
        return np.array([0, 0, 0, 0], dtype=np.float32)
    elif two_topRightY < one_bottomLeftY:  # Too far south
        return np.array([0, 0, 0, 0], dtype=np.float32)
    else:

        x_min = (
            one_bottomLeftX if one_bottomLeftX > two_bottomLeftX else two_bottomLeftX
        )
        x_max = one_topRightX if one_topRightX < two_topRightX else two_topRightX
        y_min = (
            one_bottomLeftY if one_bottomLeftY > two_bottomLeftY else two_bottomLeftY
        )
        y_max = one_topRightY if one_topRightY < two_topRightY else two_topRightY

        return np.array([x_min, x_max, y_min, y_max], dtype=np.float32)


@jit(nopython=True, nogil=True, fastmath=True, inline="always")
def get_extent(raster_transform, raster_size):
    bottomRightX = raster_transform[0] + (raster_size[0] * raster_transform[1])
    bottomRightY = raster_transform[3] + (raster_size[1] * raster_transform[5])

    return np.array(
        [raster_transform[0], bottomRightX, bottomRightY, raster_transform[3]],
        dtype=np.float32,
    )


def zonal_statistics(
    in_vector,
    output_vector=None,
    in_rasters=[],
    prefixes=[],
    stats=None,
):
    """
    ZonalStatistics.
    """
    default_stats = ["mean", "med", "std"]

    if stats is None:
        stats = default_stats

    if len(prefixes) != 0:
        if len(in_rasters) != len(prefixes):
            raise ValueError("Unable to parse prefixes.")

    if isinstance(in_rasters, list):
        if len(in_rasters) == 0:
            raise ValueError("List of rasters (in_rasters) is empty.")

    if len(stats) == 0:
        raise ValueError("Unable to parse statistics (stats).")

    # Read the raster meta:
    raster_metadata = raster_to_metadata(in_rasters[0])

    vector = None
    if output_vector is None:
        vector = open_vector(in_vector, writeable=True)
    else:
        vector = _vector_to_memory(in_vector)

    vector_metadata = _vector_to_metadata(vector)
    vector_layer = vector.GetLayer()

    # Check that projections match
    if not vector_metadata["projection_osr"].IsSame(raster_metadata["projection_osr"]):
        if output_vector is None:
            vector = reproject_vector(in_vector, in_rasters[0])
        else:
            vector_path = reproject_vector(
                in_vector, in_rasters[0], output_vector
            )
            vector = open_vector(vector_path, writeable=True)

        vector_metadata = _vector_to_metadata(vector)
        vector_layer = vector.GetLayer()

    vector_projection = vector_metadata["projection_osr"]
    raster_projection = raster_metadata["projection_osr"]

    # Read raster data in overlap
    raster_transform = np.array(raster_metadata["transform"], dtype=np.float32)
    raster_size = np.array(raster_metadata["size"], dtype=np.int32)

    raster_extent = get_extent(raster_transform, raster_size)

    vector_extent = np.array(vector_layer.GetExtent(), dtype=np.float32)
    overlap_extent = get_intersection(raster_extent, vector_extent)

    if overlap_extent is False:
        print("raster_extent: ", raster_extent)
        print("vector_extent: ", vector_extent)
        raise Exception("Vector and raster do not overlap!")

    (
        overlap_aligned_extent,
        overlap_aligned_rasterized_size,
        overlap_aligned_offset,
    ) = align_extent(raster_transform, overlap_extent, raster_size)
    overlap_transform = np.array(
        [
            overlap_aligned_extent[0],
            raster_transform[1],
            0,
            overlap_aligned_extent[3],
            0,
            raster_transform[5],
        ],
        dtype=np.float32,
    )
    overlap_size = overlap_size_calc(overlap_aligned_extent, raster_transform)

    # Loop the features
    vector_driver = ogr.GetDriverByName("Memory")
    vector_feature_count = vector_layer.GetFeatureCount()
    vector_layer.StartTransaction()

    # Create fields
    vector_layer_defn = vector_layer.GetLayerDefn()
    vector_field_counts = vector_layer_defn.GetFieldCount()
    vector_current_fields = []

    # Get current fields
    for i in range(vector_field_counts):
        vector_current_fields.append(vector_layer_defn.GetFieldDefn(i).GetName())

    # Add fields where missing
    for stat in stats:
        for i in range(len(in_rasters)):
            field_name = f"{prefixes[i]}{stat}"
            if field_name not in vector_current_fields:
                field_defn = ogr.FieldDefn(field_name, ogr.OFTReal)
                vector_layer.CreateField(field_defn)

    rasterized_features = []
    sizes = np.zeros((vector_feature_count, 4), dtype="float32")
    offsets = np.zeros((vector_feature_count, 2), dtype=np.int32)
    raster_data = None
    for raster_index, raster_value in enumerate(in_rasters):

        columns = {}
        for stat in stats:
            columns[prefixes[raster_index] + stat] = []

        fits_in_memory = True
        try:
            raster_data = raster_to_array(
                raster_value,
                pixel_offsets=[
                    overlap_aligned_offset[0],
                    overlap_aligned_offset[1],
                    overlap_aligned_rasterized_size[0],
                    overlap_aligned_rasterized_size[1],
                ],
            )
        except:
            fits_in_memory = False
            print("Raster does not fit in memory.. Doing IO for each feature.")

        for n in range(vector_feature_count):
            vector_feature = vector_layer.GetNextFeature()
            rasterized_vector = None

            if raster_index == 0:

                try:
                    vector_geom = vector_feature.GetGeometryRef()
                except:
                    vector_geom.Buffer(0)
                    Warning("Invalid geometry at : ", n)

                if vector_geom is None:
                    raise Exception("Invalid geometry. Could not fix.")

                feature_extent = vector_geom.GetEnvelope()

                # Create temp layer
                temp_vector_datasource = vector_driver.CreateDataSource(f"vector_{n}")
                temp_vector_layer = temp_vector_datasource.CreateLayer(
                    "temp_polygon", vector_projection, ogr.wkbPolygon
                )
                temp_vector_layer.CreateFeature(vector_feature.Clone())

                (
                    feature_aligned_extent,
                    feature_aligned_rasterized_size,
                    feature_aligned_offset,
                ) = align_extent(overlap_transform, feature_extent, overlap_size)

                rasterized_vector = None
                # rasterized_vector = rasterize_vector(
                #     temp_vector_layer,
                #     feature_aligned_extent,
                #     feature_aligned_rasterized_size,
                #     raster_projection,
                # )
                rasterized_features.append(rasterized_vector)

                offsets[n] = feature_aligned_offset
                sizes[n] = feature_aligned_rasterized_size

            if fits_in_memory is True:
                cropped_raster = raster_data[
                    offsets[n][1] : offsets[n][1] + int(sizes[n][1]),  # X
                    offsets[n][0] : offsets[n][0] + int(sizes[n][0]),  # Y
                ]
            else:
                cropped_raster = raster_to_array(
                    raster_value,
                    pixel_offsets=[
                        overlap_aligned_offset[0] + offsets[n][0],
                        overlap_aligned_offset[1] + offsets[n][1],
                        int(sizes[n][0]),
                        int(sizes[n][1]),
                    ],
                )

            if rasterized_features[n] is None:
                for stat in stats:
                    field_name = f"{prefixes[raster_index]}{stat}"
                    vector_feature.SetField(field_name, None)
            elif cropped_raster is None:
                for stat in stats:
                    field_name = f"{prefixes[raster_index]}{stat}"
                    vector_feature.SetField(field_name, None)
            else:
                raster_data_masked = np.ma.masked_array(
                    cropped_raster, mask=rasterized_features[n], dtype="float32"
                ).compressed()
                zonal_stats = calculate_array_stats(
                    raster_data_masked, stats
                )

                for index, stat in enumerate(stats):
                    field_name = f"{prefixes[raster_index]}{stat}"
                    vector_feature.SetField(field_name, float(zonal_stats[index]))

                vector_layer.SetFeature(vector_feature)

            progress(n, vector_feature_count, name=prefixes[raster_index])

        vector_layer.ResetReading()

    vector_layer.CommitTransaction()

    if output_vector is None:
        return vector

    return output_vector
