import datetime as dt
import os
import tempfile
from pathlib import Path

from osgeo import gdal
from osgeo import osr
from te_schemas import land_cover
from te_schemas.results import Band as JobBand

from .. import calculate_lc
from .. import utils
from .. import worker
from ..areaofinterest import AOI
from ..jobs.models import Job
from ..logger import log


def _prepare_land_cover_inputs(job: Job, area_of_interest: AOI) -> Path:
    # Select the initial and final bands from initial and final datasets
    # (in case there is more than one lc band per dataset)
    lc_initial_path = job.params["lc_initial_path"]
    lc_initial_band_index = job.params["lc_initial_band_index"]
    lc_initial_vrt = utils.save_vrt(lc_initial_path, lc_initial_band_index)
    lc_final_path = job.params["lc_final_path"]
    lc_final_band_index = job.params["lc_final_band_index"]
    lc_final_vrt = utils.save_vrt(lc_final_path, lc_final_band_index)

    # Add the lc layers to a VRT in case they don't match in resolution,
    # and set proper output bounds
    in_vrt = tempfile.NamedTemporaryFile(suffix='.vrt').name
    gdal.BuildVRT(
        in_vrt, [lc_initial_vrt, lc_final_vrt],
        resolution='highest',
        resampleAlg=gdal.GRA_NearestNeighbour,
        outputBounds=area_of_interest.
        get_aligned_output_bounds_deprecated(lc_initial_vrt),
        separate=True
    )

    return Path(in_vrt)


def compute_land_cover(
    lc_job: Job, area_of_interest: AOI, job_output_path: Path,
    dataset_output_path: Path
) -> Job:
    in_vrt = _prepare_land_cover_inputs(lc_job, area_of_interest)
    trans_matrix = land_cover.LCTransitionDefinitionDeg.Schema().loads(
        lc_job.params["trans_matrix"]
    )
    nesting = lc_job.params["legend_nesting"]

    lc_change_worker = worker.StartWorker(
        LandCoverChangeWorker,
        'calculating land cover change',
        str(in_vrt),
        str(dataset_output_path),
        trans_matrix.get_list(),
        trans_matrix.get_persistence_list(),
    )

    if lc_change_worker.success:
        lc_job.end_date = dt.datetime.now(dt.timezone.utc)
        lc_job.progress = 100
        lc_job.results.bands = [
            JobBand(
                name="Land cover (degradation)",
                metadata={
                    "year_initial":
                    lc_job.params["year_initial"],
                    "year_final":
                    lc_job.params["year_final"],
                    'trans_matrix':
                    land_cover.LCTransitionDefinitionDeg.Schema().
                    dumps(trans_matrix),
                    'nesting':
                    nesting
                },
            ),
            JobBand(
                name="Land cover (7 class)",
                metadata={
                    "year": lc_job.params["year_initial"],
                    'nesting': nesting
                }
            ),
            JobBand(
                name="Land cover (7 class)",
                metadata={
                    "year": lc_job.params["year_final"],
                    'nesting': nesting
                }
            ),
            JobBand(
                name="Land cover transitions",
                metadata={
                    "year_initial": lc_job.params["year_initial"],
                    "year_final": lc_job.params["year_final"],
                    'nesting': nesting
                }
            ),
        ]
        lc_job.results.data_path = dataset_output_path
    else:
        raise RuntimeError("Error calculating land cover change.")

    return lc_job


class LandCoverChangeWorker(worker.AbstractWorker):
    def __init__(self, in_f, out_f, trans_matrix, persistence_remap):
        worker.AbstractWorker.__init__(self)
        self.in_f = in_f
        self.out_f = out_f
        self.trans_matrix = trans_matrix
        self.persistence_remap = persistence_remap

    def work(self):
        ds_in = gdal.Open(self.in_f)

        band_initial = ds_in.GetRasterBand(1)
        band_final = ds_in.GetRasterBand(2)

        block_sizes = band_initial.GetBlockSize()
        x_block_size = block_sizes[0]
        y_block_size = block_sizes[1]
        xsize = band_initial.XSize
        ysize = band_initial.YSize

        driver = gdal.GetDriverByName("GTiff")
        ds_out = driver.Create(
            self.out_f, xsize, ysize, 4, gdal.GDT_Int16, ['COMPRESS=LZW']
        )
        src_gt = ds_in.GetGeoTransform()
        ds_out.SetGeoTransform(src_gt)
        out_srs = osr.SpatialReference()
        out_srs.ImportFromWkt(ds_in.GetProjectionRef())
        ds_out.SetProjection(out_srs.ExportToWkt())

        blocks = 0

        for y in range(0, ysize, y_block_size):
            if y + y_block_size < ysize:
                rows = y_block_size
            else:
                rows = ysize - y

            for x in range(0, xsize, x_block_size):
                if self.killed:
                    log(
                        "Processing killed by user after processing {} out of {} blocks."
                        .format(y, ysize)
                    )

                    break
                self.progress.emit(
                    100 * (float(y) + (float(x) / xsize) * y_block_size) /
                    ysize
                )

                if x + x_block_size < xsize:
                    cols = x_block_size
                else:
                    cols = xsize - x

                a_i = band_initial.ReadAsArray(x, y, cols, rows)
                a_f = band_final.ReadAsArray(x, y, cols, rows)

                a_tr = a_i * 10 + a_f
                a_tr[(a_i < 1) | (a_f < 1)] < --32768

                a_deg = a_tr.copy()

                for value, replacement in zip(
                    self.trans_matrix[0], self.trans_matrix[1]
                ):
                    a_deg[a_deg == int(value)] = int(replacement)

                # Recode transitions so that persistence classes are easier to
                # map

                for value, replacement in zip(
                    self.persistence_remap[0], self.persistence_remap[1]
                ):
                    a_tr[a_tr == int(value)] = int(replacement)

                ds_out.GetRasterBand(1).WriteArray(a_deg, x, y)
                ds_out.GetRasterBand(2).WriteArray(a_i, x, y)
                ds_out.GetRasterBand(3).WriteArray(a_f, x, y)
                ds_out.GetRasterBand(4).WriteArray(a_tr, x, y)

                blocks += 1

        if self.killed:
            os.remove(self.out_f)

            return None
        else:
            return True
