import datetime as dt
from copy import copy
from pathlib import Path

import numpy as np
import openpyxl
import qgis.core
from openpyxl.styles import Alignment, Font
from osgeo import gdal
from te_schemas.results import (
    URI,
    Band,
    DataType,
    Raster,
    RasterFileType,
    RasterResults,
)

from .. import areaofinterest, calculate, data_io, logger, summary, utils, worker
from ..jobs.models import Job


def get_main_restoration_job_params(
    combo_layer_biomass_diff: data_io.WidgetDataIOSelectTELayerExisting,
):
    band = combo_layer_biomass_diff.get_current_band()

    restoration_years = band.job.params["length_yr"]

    restoration_types = []
    for job_band in band.job.results.get_bands():
        if job_band.name == "Restoration biomass difference (tonnes CO2e per ha)":
            restoration_types.append(job_band.metadata["type"])

    return {
        "in_file_path": str(band.job.results.uri.uri),
        "restoration_years": restoration_years,
        "restoration_types": restoration_types,
        "in_file_band_infos": [
            Band.Schema().dump(b) for b in band.job.results.get_bands()
        ],
    }


def compute_biomass_restoration(
    biomass_job: Job,
    area_of_interest: areaofinterest.AOI,
    job_output_path: Path,
    dataset_output_path: Path,
    progress_callback=None,
    killed_callback=None,
) -> Job:
    #######################################################################
    # Prep files
    in_file = biomass_job.params["in_file_path"]

    # Remember the first value is an indication of whether dataset is
    # wrapped across 180th meridian
    _, wkts = area_of_interest.meridian_split("layer", "wkt", warn=False)
    bbs = area_of_interest.get_aligned_output_bounds(in_file)
    output_biomass_diff_tifs = []

    for n in range(len(wkts)):
        if len(wkts) > 1:
            output_biomass_diff_tif = (
                job_output_path.parent / f"{job_output_path.stem}_{n}.tif"
            )
        else:
            output_biomass_diff_tif = (
                job_output_path.parent / f"{job_output_path.stem}.tif"
            )
        output_biomass_diff_tifs.append(output_biomass_diff_tif)

        logger.log("Saving clipped biomass file to {}".format(output_biomass_diff_tif))
        geojson = calculate.json_geom_to_geojson(
            qgis.core.QgsGeometry.fromWkt(wkts[n]).asJson()
        )
        clip_worker = worker.StartWorker(
            calculate.ClipWorker,
            f"masking layers (part {n + 1} of {len(wkts)})",
            in_file,
            str(output_biomass_diff_tif),
            geojson,
            bbs[n],
        )

        if clip_worker.success:
            ######################################################################
            #  Calculate biomass change summary table
            logger.log("Calculating summary table...")
            rest_summary_worker = worker.StartWorker(
                RestBiomassSummaryWorker,
                f"calculating summary table (part {n + 1} of {len(wkts)})",
                str(output_biomass_diff_tif),
            )

            if rest_summary_worker.success:
                worker_result = rest_summary_worker.get_return()

                if n == 0:
                    biomass_initial, biomass_change, area_site = worker_result
                else:
                    this_initial, this_change, this_area_site = worker_result
                    biomass_initial = biomass_initial + this_initial
                    biomass_change = biomass_change + this_change
                    area_site = area_site + this_area_site
            else:
                raise RuntimeError("Error calculating biomass change summary table.")

        else:
            raise RuntimeError("Error masking input layers.")

    summary_table_output_path = job_output_path.parent / f"{job_output_path.stem}.xlsx"
    _save_summary_table(
        str(summary_table_output_path),
        biomass_initial,
        biomass_change,
        area_site,
        biomass_job.params["restoration_years"],
        biomass_job.params["restoration_types"],
    )

    # Add the biomass_dif layers to the map

    if len(output_biomass_diff_tifs) == 1:
        output_path = Path(output_biomass_diff_tifs[0])
    else:
        output_path = job_output_path.parent / f"{job_output_path.stem}.vrt"
        ds_vrt = gdal.BuildVRT(str(output_path), output_biomass_diff_tifs)
        if ds_vrt is not None:
            ds_vrt.FlushCache()
        ds_vrt = None

    # Update the band infos to use the masking value (-32767) as the file
    # no data value, so that stretches are more likely to compute correctly
    output_bands = []

    for raw_band_info in biomass_job.params["in_file_band_infos"]:
        band = Band.Schema().load(raw_band_info)
        band.no_data_value = -32767
        output_bands.append(band)

    biomass_job.results = RasterResults(
        name="biomass_restoration_summary",
        uri=URI(uri=output_path),
        rasters={
            DataType.INT16.value: Raster(
                uri=URI(uri=output_path),
                bands=output_bands,
                datatype=DataType.INT16,
                filetype=RasterFileType.GEOTIFF,
            ),
        },
    )
    biomass_job.end_date = dt.datetime.now(dt.timezone.utc)
    biomass_job.progress = 100

    return biomass_job.results


def _save_summary_table(
    out_file, biomass_initial, biomass_change, area_site, length_yr, rest_types
):
    template_summary_table_path = (
        Path(__file__).parents[1] / "data/summary_table_restoration.xlsx"
    )
    workbook = openpyxl.load_workbook(str(template_summary_table_path))

    ws_summary_sheet = workbook["Restoration Biomass Change"]
    ws_summary_sheet.cell(6, 2).value = area_site
    ws_summary_sheet.cell(7, 2).value = length_yr
    ws_summary_sheet.cell(8, 2).value = biomass_initial

    # Insert as many rows as necessary for the number of restoration types, and
    # copy the styles from the original rows, which will get pushed down

    if len(rest_types) > 1:
        offset = len(rest_types) - 1
        ws_summary_sheet.insert_rows(13, offset)

        for n in range(len(rest_types) - 1):
            _copy_style(
                ws_summary_sheet.cell(13 + offset, 1), ws_summary_sheet.cell(13 + n, 1)
            )
            _copy_style(
                ws_summary_sheet.cell(13 + offset, 2), ws_summary_sheet.cell(13 + n, 2)
            )
            _copy_style(
                ws_summary_sheet.cell(13 + offset, 3), ws_summary_sheet.cell(13 + n, 3)
            )

        # Need to remerge cells due to row insertion
        ws_summary_sheet.merge_cells(
            start_row=16 + offset, start_column=1, end_row=16 + offset, end_column=3
        )
        ws_summary_sheet.cell(16 + offset, 1).alignment = Alignment(wrap_text=True)
        ws_summary_sheet.row_dimensions[16 + offset].height = 50
        ws_summary_sheet.merge_cells(
            start_row=18 + offset, start_column=1, end_row=18 + offset, end_column=3
        )
        ws_summary_sheet.cell(18 + offset, 1).font = Font(bold=True)
        ws_summary_sheet.cell(18 + offset, 1).alignment = Alignment(wrap_text=True)
        ws_summary_sheet.row_dimensions[18 + offset].height = 60
        ws_summary_sheet.merge_cells(
            start_row=20 + offset, start_column=1, end_row=20 + offset, end_column=3
        )
        ws_summary_sheet.cell(20 + offset, 1).font = Font(bold=True)
        ws_summary_sheet.cell(20 + offset, 1).alignment = Alignment(wrap_text=True)
        ws_summary_sheet.row_dimensions[20 + offset].height = 30

    # And write the biomass differences for each restoration type

    for n in range(len(rest_types)):
        ws_summary_sheet.cell(13 + n, 1).value = rest_types[n].capitalize()
        ws_summary_sheet.cell(13 + n, 2).value = biomass_change[n]
        ws_summary_sheet.cell(13 + n, 3).value = biomass_initial + biomass_change[n]

    utils.maybe_add_image_to_sheet(
        "trends_earth_logo_bl_300width.png", ws_summary_sheet
    )

    try:
        workbook.save(out_file)
        logger.log("Summary table saved to {}".format(out_file))
    except OSError as exc:
        raise RuntimeError(
            f"Error saving output table - check that {out_file} is accessible and "
            f"not already open. - {str(exc)}"
        )


def _copy_style(a, b):
    b.font = copy(a.font)
    b.fill = copy(a.fill)
    b.border = copy(a.border)
    b.alignment = copy(a.alignment)
    b.number_format = copy(a.number_format)
    b.protection = copy(a.protection)


class RestBiomassSummaryWorker(worker.AbstractWorker):
    def __init__(self, src_file):
        worker.AbstractWorker.__init__(self)

        self.src_file = src_file

    def work(self):
        self.toggle_show_progress.emit(True)
        self.toggle_show_cancel.emit(True)

        src_ds = gdal.Open(self.src_file)

        band_biomass_initial = src_ds.GetRasterBand(1)
        # First band is initial biomass, and all other bands are for different
        # types of restoration
        n_types = src_ds.RasterCount - 1

        block_sizes = band_biomass_initial.GetBlockSize()
        xsize = band_biomass_initial.XSize
        ysize = band_biomass_initial.YSize

        x_block_size = block_sizes[0]
        y_block_size = block_sizes[1]

        src_gt = src_ds.GetGeoTransform()

        # Width of cells in longitude
        long_width = src_gt[1]
        # Set initial lat ot the top left corner latitude
        lat = src_gt[3]
        # Width of cells in latitude
        pixel_height = src_gt[5]

        area_site = 0
        biomass_initial = 0
        biomass_change = np.zeros(n_types)

        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:
                    logger.log(
                        "Processing of {} killed by user after processing {} out of {} blocks.".format(
                            self.prod_out_file, 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

                biomass_initial_array = band_biomass_initial.ReadAsArray(
                    x, y, cols, rows
                )

                # Caculate cell area for each horizontal line
                cell_areas = np.array(
                    [
                        summary.calc_cell_area(
                            lat + pixel_height * n,
                            lat + pixel_height * (n + 1),
                            long_width,
                        )
                        for n in range(rows)
                    ]
                )
                cell_areas.shape = (cell_areas.size, 1)
                # Make an array of the same size as the input arrays containing
                # the area of each cell (which is identicalfor all cells in a
                # given row - cell areas only vary among rows)
                cell_areas_array = np.repeat(cell_areas, cols, axis=1)
                # Convert cell areas to hectares
                cell_areas_array = cell_areas_array * 1e-4

                # The site area includes everything that isn't masked
                site_pixels = biomass_initial_array != -32767

                area_site = area_site + np.sum(site_pixels * cell_areas_array)
                biomass_initial = biomass_initial + np.sum(
                    site_pixels * cell_areas_array * biomass_initial_array
                )

                for n in range(n_types):
                    biomass_rest_array = src_ds.GetRasterBand(n + 2).ReadAsArray(
                        x, y, cols, rows
                    )
                    biomass_change[n] = biomass_change[n] + np.sum(
                        (biomass_rest_array) * cell_areas_array * site_pixels
                    )

                blocks += 1
            lat += pixel_height * rows
        self.progress.emit(100)

        if self.killed:
            return None

        return list((biomass_initial, biomass_change, area_site))
