from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import ee
from te_schemas.schemas import BandInfo

from .util import TEImage


def soc(
    year_initial, year_final, fl, trans_matrix, nesting, dl_annual_lc, logger
):
    """Calculate SOC indicator."""
    logger.debug("Entering soc function.")

    # soc
    soc = ee.Image("users/geflanddegradation/toolbox_datasets/soc_sgrid_30cm")
    soc_t0 = soc.updateMask(soc.neq(-32768))

    # Reference all SOC calculations to the year 2000
    soc_t0_year = 2000

    # First band in the LC layer stack is 1992
    lc_band0_year = 1992

    # land cover - note it needs to be reprojected to match soc so that it can
    # be output to cloud storage in the same stack
    lc = ee.Image(
        "users/geflanddegradation/toolbox_datasets/lcov_esacc_1992_2020"
    ).select(
        ee.List.sequence(
            soc_t0_year - lc_band0_year, year_final - lc_band0_year, 1
        )
    ).reproject(crs=soc.projection())
    lc = lc.where(lc.eq(9999), -32768)
    lc = lc.updateMask(lc.neq(-32768))

    if fl == 'per pixel':
        # Setup a raster of climate regimes to use for coding Fl automatically
        climate = (
            ee.Image("users/geflanddegradation/toolbox_datasets/ipcc_climate_zones")
            .remap([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
                   [0, 2, 1, 2, 1, 2, 1, 2, 1, 5, 4, 4, 3])
        )  # yapf: disable
        clim_fl = climate.remap(
            [0, 1, 2, 3, 4, 5],
            [0, 0.8, 0.69, 0.58, 0.48, 0.64])  # yapf: disable
    # create empty stacks to store annual land cover maps
    stack_lc = ee.Image().select()

    # create empty stacks to store annual soc maps
    stack_soc = ee.Image().select()

    # loop through all the years in the period of analysis to compute changes in SOC

    for k in range(year_final - soc_t0_year):
        # land cover map reclassified to UNCCD 7 classes (1: forest, 2:
        # grassland, 3: cropland, 4: wetland, 5: artifitial, 6: bare, 7: water)
        lc_t0 = lc.select(k).remap(
            nesting.get_list()[0],
            nesting.get_list()[1]
        )

        lc_t1 = lc.select(k + 1).remap(
            nesting.get_list()[0],
            nesting.get_list()[1]
        )

        if (k == 0):
            # compute transition map (first digit for baseline land cover, and
            # second digit for target year land cover)
            lc_tr = lc_t0.multiply(trans_matrix.get_multiplier()).add(lc_t1)

            # compute raster to register years since transition
            tr_time = ee.Image(2).where(lc_t0.neq(lc_t1), 1)
        else:
            # Update time since last transition. Add 1 if land cover remains
            # constant, and reset to 1 if land cover changed.
            tr_time = tr_time.where(lc_t0.eq(lc_t1), tr_time.add(
                ee.Image(1)
            )).where(lc_t0.neq(lc_t1), ee.Image(1))

            # compute transition map (first digit for baseline land cover, and
            # second digit for target year land cover), but only update where
            # changes actually ocurred.
            lc_tr_temp = lc_t0.multiply(10).add(lc_t1)
            lc_tr = lc_tr.where(lc_t0.neq(lc_t1), lc_tr_temp)

        # stock change factor for land use - note the 99 and -99 will be
        # recoded using the chosen Fl option
        lc_tr_fl_0 = lc_tr.remap(
            [11, 12, 13, 14, 15, 16, 17,
             21, 22, 23, 24, 25, 26, 27,
             31, 32, 33, 34, 35, 36, 37,
             41, 42, 43, 44, 45, 46, 47,
             51, 52, 53, 54, 55, 56, 57,
             61, 62, 63, 64, 65, 66, 67,
             71, 72, 73, 74, 75, 76, 77],
            [1,     1,    99,        1, 0.1, 0.1, 1,
             1,     1,    99,        1, 0.1, 0.1, 1,
             -99, -99,     1, 1 / 0.71, 0.1, 0.1, 1,
             1,      1, 0.71,        1, 0.1, 0.1, 1,
             2,      2,    2,        2,   1,   1, 1,
             2,      2,    2,        2,   1,   1, 1,
             1,      1,    1,        1,   1,   1, 1])  # yapf: disable

        if fl == 'per pixel':
            lc_tr_fl = lc_tr_fl_0.where(lc_tr_fl_0.eq(99), clim_fl).where(
                lc_tr_fl_0.eq(-99),
                ee.Image(1).divide(clim_fl)
            )
        else:
            lc_tr_fl = lc_tr_fl_0.where(lc_tr_fl_0.eq(99), fl).where(
                lc_tr_fl_0.eq(-99),
                ee.Image(1).divide(fl)
            )

        # stock change factor for management regime
        lc_tr_fm = lc_tr.remap([11, 12, 13, 14, 15, 16, 17,
                                21, 22, 23, 24, 25, 26, 27,
                                31, 32, 33, 34, 35, 36, 37,
                                41, 42, 43, 44, 45, 46, 47,
                                51, 52, 53, 54, 55, 56, 57,
                                61, 62, 63, 64, 65, 66, 67,
                                71, 72, 73, 74, 75, 76, 77],
                               [1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1])  # yapf: disable

        # stock change factor for input of organic matter
        lc_tr_fo = lc_tr.remap([11, 12, 13, 14, 15, 16, 17,
                                21, 22, 23, 24, 25, 26, 27,
                                31, 32, 33, 34, 35, 36, 37,
                                41, 42, 43, 44, 45, 46, 47,
                                51, 52, 53, 54, 55, 56, 57,
                                61, 62, 63, 64, 65, 66, 67,
                                71, 72, 73, 74, 75, 76, 77],
                               [1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1,
                                1, 1, 1, 1, 1, 1, 1])  # yapf: disable

        if (k == 0):
            soc_chg = (
                soc_t0.subtract(
                    (
                        soc_t0.multiply(lc_tr_fl).multiply(lc_tr_fm).
                        multiply(lc_tr_fo)
                    )
                )
            ).divide(20)

            # compute final SOC stock for the period
            soc_t1 = soc_t0.subtract(soc_chg)

            # add to land cover and soc to stacks from both dates for the first
            # period
            stack_lc = stack_lc.addBands(lc_t0).addBands(lc_t1)
            stack_soc = stack_soc.addBands(soc_t0).addBands(soc_t1)

        else:
            # compute annual change in soc (updates from previous period based
            # on transition and time <20 years)
            soc_chg = (
                soc_chg.where(
                    lc_t0.neq(lc_t1), (
                        stack_soc.select(k).subtract(
                            stack_soc.select(k).multiply(lc_tr_fl).
                            multiply(lc_tr_fm).multiply(lc_tr_fo)
                        )
                    ).divide(20)
                ).where(tr_time.gt(20), 0)
            )

            # compute final SOC for the period
            socn = stack_soc.select(k).subtract(soc_chg)

            # add land cover and soc to stacks only for the last year in the
            # period
            stack_lc = stack_lc.addBands(lc_t1)
            stack_soc = stack_soc.addBands(socn)

    # compute soc percent change for the analysis period
    soc_pch = (
        (
            (
                stack_soc.select(year_final - soc_t0_year).subtract(
                    stack_soc.select(year_initial - soc_t0_year)
                )
            ).divide(stack_soc.select(year_initial - soc_t0_year))
        ).multiply(100)
    ).rename(f'Percent_SOC_increase_{year_initial}-{year_final}')

    logger.debug("Setting up output.")
    out = TEImage(
        soc_pch, [
            BandInfo(
                "Soil organic carbon (degradation)",
                add_to_map=True,
                metadata={
                    'year_initial': year_initial,
                    'year_final': year_final,
                    'trans_matrix': trans_matrix.dumps(),
                    'nesting': nesting.dumps()
                }
            )
        ]
    )

    logger.debug("Adding annual SOC layers.")

    # Setup a list of the years included
    years = [*range(year_initial, year_final + 1)]

    # Output annual SOC layers
    d_soc = []

    for year in years:
        if year == years[0]:
            soc_stack_out = stack_soc.select(year - soc_t0_year)
        else:
            soc_stack_out = soc_stack_out.addBands(
                stack_soc.select(year - soc_t0_year)
            )

        if (year == year_initial) or (year == year_final):
            add_to_map = True
        else:
            add_to_map = False
        d_soc.append(
            BandInfo(
                "Soil organic carbon",
                add_to_map=add_to_map,
                metadata={'year': year}
            )
        )
    soc_stack_out = soc_stack_out.rename([f'SOC_{year}' for year in years])
    out.addBands(soc_stack_out, d_soc)

    if dl_annual_lc:
        logger.debug("Adding all annual LC layers.")
        d_lc = []

        for year in years:
            if year == years[0]:
                lc_stack_out = stack_lc.select(year - soc_t0_year)
            else:
                lc_stack_out = lc_stack_out.addBands(
                    stack_lc.select(year - soc_t0_year)
                )
            d_lc.append(
                BandInfo(
                    "Land cover (7 class)",
                    metadata={
                        'year': year,
                        'nesting': nesting.dumps()
                    }
                )
            )
        lc_stack_out = lc_stack_out.rename(
            [f'Land_cover_{year}' for year in years]
        )
        out.addBands(lc_stack_out, d_lc)
    else:
        logger.debug("Adding initial and final LC layers.")
        lc_initial = stack_lc.select(year_initial - soc_t0_year
                                     ).rename(f'Land_cover_{year_initial}')
        lc_final = stack_lc.select(year_final - soc_t0_year
                                   ).rename(f'Land_cover_{year_final}')
        out.addBands(
            lc_initial.addBands(lc_final), [
                BandInfo(
                    "Land cover (7 class)",
                    metadata={
                        'year': year_initial,
                        'nesting': nesting.dumps()
                    }
                ),
                BandInfo(
                    "Land cover (7 class)",
                    metadata={
                        'year': year_final,
                        'nesting': nesting.dumps()
                    }
                )
            ]
        )

    out.image = out.image.unmask(-32768).int16()

    return out
