"""remotepixel.cbers_ndvi module."""

import re
from functools import partial
from concurrent import futures

import numpy as np
import numexpr as ne

import rasterio
from rasterio import warp

from remotepixel.utils import cbers_parse_scene_id, get_area
from rio_tiler.utils import linear_rescale, array_to_img, get_colormap, b64_encode_img

np.seterr(divide="ignore", invalid="ignore")

CBERS_BUCKET = "s3://cbers-pds"


def point(scene, coordinates, expression):
    """Point handler."""
    bands = tuple(set(re.findall(r"b(?P<bands>[0-9]{1,2})", expression)))

    scene_params = cbers_parse_scene_id(scene)
    cbers_address = f'{CBERS_BUCKET}/{scene_params["key"]}'
    addresses = [
        "{}/{}_BAND{}.tif".format(cbers_address, scene, band) for band in bands
    ]

    def worker(address):
        """Point worker."""
        with rasterio.open(address) as band:
            lon_srs, lat_srs = warp.transform(
                "EPSG:4326", band.crs, [coordinates[0]], [coordinates[1]]
            )
            point = list(band.sample([(lon_srs[0], lat_srs[0])]))[0]
        return point[0]

    try:
        with futures.ThreadPoolExecutor(max_workers=3) as executor:
            data = list(executor.map(worker, addresses))

            ctx = {}
            for bdx, b in enumerate(bands):
                ctx["b{}".format(b)] = data[bdx]
            ratio = np.nan_to_num(ne.evaluate(expression, local_dict=ctx))
    except IndexError:
        ratio = 0.0

    date = (
        scene_params["acquisitionYear"]
        + "-"
        + scene_params["acquisitionMonth"]
        + "-"
        + scene_params["acquisitionDay"]
    )

    return {"ndvi": ratio, "scene": scene, "date": date}


def area(
    scene,
    bbox,
    expression,
    expression_range=[-1, 1],
    bbox_crs="epsg:4326",
    out_crs="epsg:3857",
):
    """Area handler."""
    max_img_size = 512

    bands = tuple(set(re.findall(r"b(?P<bands>[0-9]{1,2})", expression)))

    scene_params = cbers_parse_scene_id(scene)
    cbers_address = f'{CBERS_BUCKET}/{scene_params["key"]}'
    addresses = [
        "{}/{}_BAND{}.tif".format(cbers_address, scene, band) for band in bands
    ]

    _worker = partial(
        get_area,
        bbox=bbox,
        max_img_size=max_img_size,
        bbox_crs=bbox_crs,
        out_crs=out_crs,
    )
    with futures.ThreadPoolExecutor(max_workers=3) as executor:
        data = np.concatenate(list(executor.map(_worker, addresses)))
        if not np.any(data):
            raise Exception("No valid data in array")

        mask = np.all(data != 0, axis=0).astype(np.uint8) * 255

        ctx = {}
        for bdx, b in enumerate(bands):
            ctx["b{}".format(b)] = data[bdx]
        ratio = np.nan_to_num(ne.evaluate(expression, local_dict=ctx))

    ratio = np.where(
        mask, linear_rescale(ratio, in_range=expression_range, out_range=[0, 255]), 0
    ).astype(np.uint8)

    img = array_to_img(ratio, mask, get_colormap(name="cfastie"))
    ndvi = b64_encode_img(img, "jpeg")

    date = (
        scene_params["acquisitionYear"]
        + "-"
        + scene_params["acquisitionMonth"]
        + "-"
        + scene_params["acquisitionDay"]
    )

    return {"ndvi": ndvi, "scene": scene, "date": date}
