import numpy as np
import warnings
from pathlib import Path
import datetime
from cyclesgym.managers.operation import OperationManager
from cyclesgym.envs.utils import date2ydoy

__all__ = ['Fertilizer', 'FixedRateNFertilizer']

class Implementer(object):

    def year2opyear(self, year: int) -> int:
        """
        Convert normal year to years used in operation.

        Cycles and cyclesgym use normal years (e.g. 1980) but, in the operation
        file, they start counting from the first simulation year, i.e., if the
        sim start in 1980, this is indicated as year 1 in the operation file.
        """
        return year - self.start_year + 1

    def _is_new_action(self, year: int, doy: int, *args, **kwargs) -> bool:
        raise NotImplementedError

    def reset(self) -> bool:
        raise NotImplementedError

    def _check_valid_action(self, action_details: dict):
        raise NotImplementedError

    def _update_operation(self, year, doy, operation, *args, **kwargs):
        raise NotImplementedError

    def _check_collision(self, year, doy, operation_details):
        raise NotImplementedError

    def _get_operation_key(self, year, doy, operation_details):
        raise NotImplementedError

    def _create_new_operation(self, year, key, operation_details):
        raise NotImplementedError

    def implement_action(self, date: datetime.date, operation_details: dict) -> bool:
        """
        Write new operation file.

        We check if the new action is different from the previously specified
        one. If it is, we overwrite the operation file and we indicate cycles
        should be rerun. Othwerwise, we do nothing.

        Parameters
        ----------
        date: datetime.date
            Date
        operation_details: dict
                    details of the operation to be implemented

        Returns
        -------
        rerun_cycles: bool
            if true, it means the operation file has been changed and cycles
            should be rerun.
        """
        # Check all and only affected crops are specified
        self._check_valid_action(operation_details)

        year, doy = date2ydoy(date)
        year = self.year2opyear(year)

        if self._is_new_action(year, doy, operation_details):
            # Check for collision
            operation, collision = self._check_collision(year, doy, operation_details)
            if collision:
                # Take existing operations and update masses
                op = self._update_operation(year, doy, operation, operation_details)
            else:
                op = self._create_new_operation(year, doy, operation_details)

            # Insert and write operation
            self.operation_manager.insert_new_operations(op, force=True)
            self.operation_manager.save(self.operation_fname)

            rerun_cycles = True
        else:
            # If not new action, don't update operation and don't rerun
            rerun_cycles = False
        return rerun_cycles

class Fertilizer(Implementer):
    valid_nutrients = ['C_Organic', 'C_Charcoal', 'N_Organic', 'N_Charcoal',
                       'N_NH4', 'N_NO3', 'P_Organic', 'P_CHARCOAL',
                       'P_INORGANIC', 'K', 'S']

    def __init__(self,
                 operation_manager: OperationManager,
                 operation_fname: Path,
                 affected_nutrients: list,
                 start_year: int) -> None:
        self.operation_manager = operation_manager
        self.operation_fname = operation_fname
        assert all([nutrient in self.valid_nutrients for nutrient in
                    affected_nutrients]), f'Affected nutrients should be in ' \
                                          f'{self.valid_nutrients}. Got' \
                                          f'{affected_nutrients} instead.'
        self.affected_nutrients = affected_nutrients
        self.start_year = start_year

    def _is_new_action(self, year: int, doy: int, new_masses: dict) -> bool:
        """
        Is the fertilization recommended for (year, doy) different from previous one?

        Parameters
        ----------
        year: int
            Year
        doy: int
            Day of the year
        new_masses: dict
            keys are nutrients and values are masses in Kg

        Returns
        -------
        new: bool
        """
        # Get old operation
        old_op = self.operation_manager.op_dict.get(
            (year, doy, 'FIXED_FERTILIZATION'))
        if old_op is None:
            old_masses = {n: 0 for n in self.affected_nutrients}
        else:
            total_mass = old_op['MASS']
            old_masses = {n: v * total_mass for n, v in old_op.items()
                          if n in self.affected_nutrients}
        return old_masses != new_masses

    def _get_operation_key(self, year, doy):
        return (year, doy, 'FIXED_FERTILIZATION')

    def _check_valid_action(self, action_details: dict):
        assert all(n in self.affected_nutrients for n in
                   action_details.keys()), f'You can only specify fertilization masses for {self.affected_nutrients}'
        assert all(n in action_details.keys() for n in
                   self.affected_nutrients), f'You must specify fertilization masses for all {self.affected_nutrients}'

    def _create_new_operation(self, year, doy, masses: dict) -> dict:
        key = self._get_operation_key(year, doy)
        op_val = {
            'SOURCE': 'Unknown',
            'MASS': 0,
            'FORM': 'Unknown',
            'METHOD': 'Unknown',
            'LAYER': 1,  # TODO: Handle different layers
        }

        # Initialize all nutrients to zero
        op_val.update({n: 0 for n in self.valid_nutrients})

        # Update value for affected nutrients
        total_mass = sum([m for m in masses.values()])
        op_val.update({'MASS': total_mass})
        op_val.update({nutrient: mass / total_mass for
                       (nutrient, mass) in masses.items()})
        op = {key: op_val}
        return op

    def reset(self) -> bool:
        """
        Set to zero the masses for all the affected nutrients.

        Before starting a new simulation, we want to make sure that the only
        way the affected nutrients can be provided is through our
        fertilization. This is why we reset everything to zero (overwriting
        during the unrolling of the env may not be sufficient as there may be
        fertilization event happening between t and t + delta that we could not
        overwrite.
        """
        doy = 1
        rerun_cycles = False
        for op_k, op_v in self.operation_manager.op_dict.items():
            if op_k[-1] == 'FIXED_FERTILIZATION':
                start_year = op_k[0]
                doy = op_k[1]
                new_op = self._update_operation(start_year, doy,
                                                old_op = op_v,
                                                new_masses={n: 0 for n in self.affected_nutrients},
                                                mode='absolute')
                self.operation_manager.insert_new_operations(
                    new_op, force=True)
                rerun_cycles = True
        self.operation_manager.save(self.operation_fname)
        return rerun_cycles

    def _check_collision(self, year, doy, operation_details):
        key = (year, doy, 'FERTILIZATION')
        operation = self.operation_manager.op_dict.get(key)
        collision = operation is not None
        return operation, collision

    def _update_operation(self, year, doy,
                          old_op: dict, new_masses: dict,
                          mode: str = 'absolute') -> dict:
        """
        Update the masses of a fertilization operation.

        Parameters
        ----------
        old_op: dict
            Dictionary of fertilization operation to update
        new_masses: dict
            Dict of new masses. Keys should be possible nutrients
        mode: str
            With 'increment' the mass is added to the one already present
            in the operation. With 'absolute' we use the new mass regardless of
            what was already there
        Returns
        -------
        new_op: dict
            Dictionary of updated fertilization operation
        """

        # TODO: include mode in the call of the Abstract class in some way. Right now abstract class has no "mode"

        key = self._get_operation_key(year, doy) # TODO: Weird we get a dict with massese as input and return a dict {k: masses} as output. Can't we add the key later?

        assert mode in ['increment', 'absolute'], \
            'Can only update in increment or absolute mode'

        assert all(k in self.valid_nutrients for k in new_masses.keys()), \
            f'New masses can only specify valid nutrients ' \
            f'({self.valid_nutrients}) as keys'

        # Copy operation
        new_op = old_op.copy()

        # Init final masses
        final_masses = {k: 0 for k in self.valid_nutrients}

        for i, n in enumerate(self.valid_nutrients):
            # Update with new masses when specified
            if n in new_masses.keys():
                final_masses[n] = new_masses[n] if mode == 'absolute' else \
                    new_masses[n] + old_op[n]
            # Copy old masses otherwise
            else:
                final_masses[n] = old_op[n] * old_op['MASS']

        # Update dict
        total_mass = sum(final_masses.values())
        new_op['MASS'] = total_mass
        if total_mass > 0:
            new_op.update(
                {nutrient: final_mass / total_mass for (nutrient, final_mass)
                 in final_masses.items()})
        else:
            new_op.update({nutrient: 0 for nutrient in self.valid_nutrients})

        return {key: new_op}

class FixedRateNFertilizer(Fertilizer):
    def __init__(self, operation_manager: OperationManager,
                 operation_fname: Path,
                 rate: float,
                 start_year: int):
        nutrients = ['N_NH4', 'N_NO3']
        super(FixedRateNFertilizer, self).__init__(operation_manager,
                                                   operation_fname,
                                                   nutrients,
                                                   start_year)
        assert 0 <= rate <= 1, f'Rate must be in [0, 1]. It is {rate} instead'
        self.rate = rate

    def convert_mass(self, mass):
        masses = {'N_NH4': mass * self.rate,
                  'N_NO3': mass * (1 - self.rate)}
        return masses

    # TODO: This should have same signature as parent method!
    def implement_action(self, date: datetime.date, mass: float):
        masses = self.convert_mass(mass)
        return super(FixedRateNFertilizer, self).implement_action(date, masses)

class Planter(Implementer):
    valid_crops = ['CornRM.90', 'CornRM.100', 'CornRM.110', 'CornSilageRM.90', 'CornSilageRM.100', 'CornSilageRM.110',
                   'SorghumSS', 'SorghumMS', 'SorghumLS', 'SweetSorghum', 'SoybeanMG.3', 'SoybeanMG.4', 'SoybeanMG.5',
                   'OatsGrain', 'OatsHay', 'RygrassAnnualGrazing', 'SpringWheat', 'SpringBarley', 'WinterWheat',
                   'WinterBarley', 'Chickpea', 'SpringPeas', 'WinterPeas', 'SpringLentils', 'WinterLentils',
                   'SpringCanola', 'WinterCanola', 'Potatoes', 'Millet', 'Sesame_Beta', 'Teff_Beta', 'Alfalfa',
                   'LotusCorniculatus', 'WhiteClover', 'Orchardgrass', 'TallFescueGrazing', 'Switchgrass', 'Miscanthus',
                   'Willow', 'C3_weed', 'C4_weed']

    def __init__(self,
                 operation_manager: OperationManager,
                 operation_fname: Path,
                 affected_crops: list,
                 start_year: int) -> None:
        self.operation_manager = operation_manager
        self.operation_fname = operation_fname
        assert all([crop in self.valid_crops for crop in
                    affected_crops]), f'Affected crops should be in ' \
                                          f'{self.valid_crops}. Got' \
                                          f'{affected_crops} instead.'
        self.affected_crops = affected_crops
        self.start_year = start_year

    def _is_new_action(self, year: int, doy: int, crop_details: dict) -> bool:
        """
        Is the crop recommended for (year, doy) different from previous one?

        Parameters
        ----------
        year: int
            Year
        doy: int
            Day of the year
        crop: str
            crop to be planted

        Returns
        -------
        new: bool
        """
        # Get old operation
        old_op = self.operation_manager.op_dict.get(
            (year, doy, 'PLANTING'))
        if old_op is None:
            return True
        else:
            return old_op['CROP'] == crop_details['CROP']

    def _check_valid_action(self, action_details: dict):
        assert (action_details['CROP'] in self.valid_crops), f'You can only specify planting date for ' \
                                                             f'{self.valid_crops}'

    def reset(self) -> bool:
        """
        Set to zero the planting events of all the affected crops.

        Before starting a new simulation, we want to make sure that the only
        way the affected crops can be planted is using the planter implementer.
        This is why we delete all the operations that involve planting affected
        crops
        """
        rerun_cycles = True
        op_to_delete = []
        for op_k, op_v in self.operation_manager.op_dict.items():
            if op_k[-1] == 'PLANTING' and op_v['CROP'] in self.affected_crops:
                op_to_delete.append(op_k)

        self.operation_manager.delete_operations(op_to_delete)
        self.operation_manager.save(self.operation_fname)
        # TODO: Understand when we can avoid rerunning cycles
        return rerun_cycles

    def _get_operation_key(self, year, doy):
        return (year, doy, 'PLANTING')

    def _check_collision(self, year, doy, operation_details):
        planting_events = [key for key in self.operation_manager.op_dict.keys() if key[2] == 'PLANTING']
        operation = [key for key in planting_events if key[0] == year]
        collision = False
        if len(operation) > 0:
            operation = operation[0]
            operation = {operation: self.operation_manager.op_dict[operation]}
            collision = True
        return operation, collision

    def _update_operation(self, year, doy, old_operation, operation_details) -> dict:
        """
        Update the planting operation.

        Parameters
        ----------
        old_operation: dict
            Dictionary of operation detials of old op
        operation_details: dict
            Dict of new operation.
        -------
        new_op: dict
            Dictionary of updated planting operation
        """

        assert operation_details['CROP'] in self.valid_crops, \
            f'New planting operation can only specify valid crops \n' \
            f'({self.valid_crops}) as planted crop \n' \
            f"Given {operation_details['CROP']}"

        doy = operation_details.pop('DOY')
        for item in old_operation.values():
            item.update(operation_details)

        old_key = list(old_operation.keys())[0]
        new_key = (old_key[0], doy, old_key[2])
        old_operation[new_key] = old_operation.pop(old_key)

        self.operation_manager.op_dict.pop(old_key)
        return old_operation

    def _create_new_operation(self, year, doy, operation_details: dict) -> dict:
        key = self._get_operation_key(year, doy)
        op_val = {'DOY': 1,
                  'END_DOY': -999,
                  'MAX_SMC': -999,
                  'MIN_SMC': 0.0,
                  'MIN_SOIL_TEMP': 0.0,
                  'CROP': None,
                  'USE_AUTO_IRR': 0,
                  'USE_AUTO_FERT': 0,
                  'FRACTION': 1.0,
                  'CLIPPING_START': 1,
                  'CLIPPING_END': 366
                  }

        # Initialize all nutrients to zero
        op_val.update(operation_details)
        key = (key[0], op_val.pop('DOY'), key[2])
        return {key: op_val}

class RotationPlanter(Planter):
    def __init__(self, operation_manager: OperationManager,
                 operation_fname: Path,
                 rotation_crops: list,
                 start_year: int):
        super(RotationPlanter, self).__init__(operation_manager,
                                              operation_fname,
                                              rotation_crops,
                                              start_year)

    def convert_action_to_dict(self, crop_categorical: int, doy: int,
                               end_doy: int, max_smc: int):
        """crop_categorical: integer representing the crop in the crop rotation
        doy: integer lower bound of the planting date, as number of weeks starting from the
            first of march
        end_doy: upper bound of the planting date, as number of weeks from the doy
        max_smc: maximum moisture for automated planting"""

        doy = 90 + doy * 7
        end_doy = doy + end_doy * 7
        max_smc = float((1 + max_smc) / 10.)

        operation_det = {'DOY': doy,
                         'CROP': self.affected_crops[crop_categorical],
                         'END_DOY': end_doy,
                         'MAX_SMC': max_smc}
        return operation_det

    def implement_action(self, date: datetime.date, crop_categorical: int, doy: int,
                         END_DOY: int, MAX_SMC: int):
        operation_det = self.convert_action_to_dict(crop_categorical, doy, END_DOY, MAX_SMC)
        return super(RotationPlanter, self).implement_action(date, operation_det)

class RotationPlanterFixedPlanting(RotationPlanter):

    def implement_action(self, date: datetime.date, crop_categorical: int, doy: int):
        operation_det = self.convert_action_to_dict(crop_categorical, doy, 0, 0)
        operation_det['END_DOY'] = -999
        operation_det['MAX_SMC'] = -999
        return super(RotationPlanter, self).implement_action(date, operation_det)


import os
from unittest import mock

import poptorch
import pytest
import torch
from lightning_utilities.core.imports import package_available
from torch.utils.data import DistributedSampler

if package_available("lightning"):
    from lightning.pytorch import Callback, Trainer, seed_everything
    from lightning.pytorch.core.module import LightningModule
    from lightning.pytorch.demos.boring_classes import BoringModel
    from lightning.pytorch.trainer.states import RunningStage, TrainerFn
    from lightning.pytorch.utilities.exceptions import MisconfigurationException

elif package_available("pytorch_lightning"):
    from pytorch_lightning import Callback, Trainer, seed_everything
    from pytorch_lightning.core.module import LightningModule
    from pytorch_lightning.demos.boring_classes import BoringModel
    from pytorch_lightning.trainer.states import RunningStage, TrainerFn
    from pytorch_lightning.utilities.exceptions import MisconfigurationException

from lightning_graphcore import IPUStrategy
from lightning_graphcore.accelerator import IPUAccelerator
from lightning_graphcore.precision import IPUPrecision

from tests.helpers import ClassifDataModule, IPUClassificationModel, IPUModel

def test_auto_device_count():
    assert IPUAccelerator.auto_device_count() == 4

@mock.patch("lightning_graphcore.accelerator.IPUAccelerator.is_available", return_value=False)
def test_fail_if_no_ipus(_, tmpdir):  # noqa: PT019
    with pytest.raises(
        MisconfigurationException,
        match="`IPUAccelerator` can not run on your system since the accelerator is not available.",
    ):
        Trainer(default_root_dir=tmpdir, accelerator=IPUAccelerator(), devices=1)

@pytest.mark.xfail()  # todo
def test_accelerator_selected(tmpdir):
    assert IPUAccelerator.is_available()
    trainer = Trainer(default_root_dir=tmpdir, accelerator="ipu", devices=1)
    assert isinstance(trainer.accelerator, IPUAccelerator)

def test_warning_if_ipus_not_used():
    with pytest.warns(UserWarning, match="IPU available but not used. Set `accelerator` and `devices`"):
        Trainer(accelerator="cpu")

def test_no_warning_strategy(tmpdir):
    with pytest.warns(None) as record:
        Trainer(default_root_dir=tmpdir, max_epochs=1, strategy=IPUStrategy(training_opts=poptorch.Options()))
    assert len(record) == 0

@pytest.mark.parametrize(
    "devices",
    [1, 4],
)
def test_all_stages(tmpdir, devices):
    model = IPUModel()
    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, strategy=IPUStrategy(), devices=devices)
    trainer.fit(model)
    trainer.validate(model)
    trainer.test(model)
    trainer.predict(model)

@pytest.mark.parametrize(
    "devices",
    [
        1,
        pytest.param(
            4,
            marks=pytest.mark.xfail(  # fixme
                AssertionError, reason="Invalid batch dimension: In the input torch.Size([1, 32]), ..."
            ),
        ),
    ],
)
def test_inference_only(tmpdir, devices):
    model = IPUModel()

    trainer = Trainer(
        default_root_dir=tmpdir, fast_dev_run=True, strategy=IPUStrategy(accelerator=IPUAccelerator()), devices=devices
    )
    trainer.validate(model)
    trainer.test(model)
    trainer.predict(model)

@pytest.mark.xfail(reason="TODO: issues with latest poptorch")
@mock.patch.dict(os.environ, os.environ.copy(), clear=True)
def test_optimization(tmpdir):
    seed_everything(42)

    dm = ClassifDataModule(length=1024)
    model = IPUClassificationModel()

    trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, accelerator=IPUAccelerator(), devices=2)

    # fit model
    trainer.fit(model, dm)
    assert trainer.state.finished, f"Training failed with {trainer.state}"
    assert dm.trainer is not None

    # validate
    result = trainer.validate(datamodule=dm)
    assert dm.trainer is not None
    assert result[0]["val_acc"] > 0.7

    # test
    result = trainer.test(model, datamodule=dm)
    assert dm.trainer is not None
    test_result = result[0]["test_acc"]
    assert test_result > 0.6

    # test saved model
    model_path = os.path.join(tmpdir, "model.pt")
    trainer.save_checkpoint(model_path)

    model = IPUClassificationModel.load_from_checkpoint(model_path)

    trainer = Trainer(default_root_dir=tmpdir, accelerator=IPUAccelerator(), devices=2)

    result = trainer.test(model, datamodule=dm)
    saved_result = result[0]["test_acc"]
    assert saved_result == test_result

def test_half_precision(tmpdir):
    class TestCallback(Callback):
        def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
            assert trainer.precision == "16-mixed"
            raise SystemExit

    model = IPUModel()
    trainer = Trainer(
        default_root_dir=tmpdir,
        fast_dev_run=True,
        strategy=IPUStrategy(accelerator=IPUAccelerator(), precision_plugin=IPUPrecision("16-mixed")),
        devices=1,
        callbacks=TestCallback(),
    )
    assert isinstance(trainer.strategy.precision_plugin, IPUPrecision)
    assert trainer.strategy.precision_plugin.precision == "16-mixed"
    with pytest.raises(SystemExit):
        trainer.fit(model)

def test_pure_half_precision(tmpdir):
    class TestCallback(Callback):
        def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
            assert trainer.strategy.precision_plugin.precision == "16-mixed"
            for param in trainer.strategy.model.parameters():
                assert param.dtype == torch.float16
            raise SystemExit

    model = IPUModel()
    model = model.half()
    trainer = Trainer(
        default_root_dir=tmpdir,
        fast_dev_run=True,
        strategy=IPUStrategy(accelerator=IPUAccelerator(), precision_plugin=IPUPrecision("16-mixed")),
        devices=1,
        callbacks=TestCallback(),
    )

    assert isinstance(trainer.strategy, IPUStrategy)
    assert isinstance(trainer.strategy.precision_plugin, IPUPrecision)
    assert trainer.strategy.precision_plugin.precision == "16-mixed"

    changed_dtypes = [torch.float, torch.float64]
    data = [torch.zeros((1), dtype=dtype) for dtype in changed_dtypes]
    new_data = trainer.strategy.batch_to_device(data)
    assert all(val.dtype is torch.half for val in new_data), "".join(
        [f"{dtype}: {val.dtype}" for dtype, val in zip(changed_dtypes, new_data)]
    )

    not_changed_dtypes = [torch.uint8, torch.int8, torch.int32, torch.int64]
    data = [torch.zeros((1), dtype=dtype) for dtype in not_changed_dtypes]
    new_data = trainer.strategy.batch_to_device(data)
    assert all(val.dtype is dtype for val, dtype in zip(new_data, not_changed_dtypes)), "".join(
        [f"{dtype}: {val.dtype}" for dtype, val in zip(not_changed_dtypes, new_data)]
    )

    with pytest.raises(SystemExit):
        trainer.fit(model)

def test_true_half_precision(tmpdir):
    class TestCallback(Callback):
        def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
            assert trainer.strategy.precision_plugin.precision == "16-true"
            for param in trainer.strategy.model.parameters():
                assert param.dtype == torch.float16
            raise SystemExit

    model = IPUModel()
    trainer = Trainer(
        default_root_dir=tmpdir,
        fast_dev_run=True,
        strategy=IPUStrategy(accelerator=IPUAccelerator(), precision_plugin=IPUPrecision("16-true")),
        devices=1,
        callbacks=TestCallback(),
    )

    assert isinstance(trainer.strategy, IPUStrategy)
    assert isinstance(trainer.strategy.precision_plugin, IPUPrecision)
    assert trainer.strategy.precision_plugin.precision == "16-true"

    # TODO uncomment
    # changed_dtypes = [torch.float, torch.float64]
    # data = [torch.zeros((1), dtype=dtype) for dtype in changed_dtypes]
    # new_data = trainer.strategy.batch_to_device(data)
    # assert all(val.dtype is torch.half for val in new_data), "".join(
    #     [f"{dtype}: {val.dtype}" for dtype, val in zip(changed_dtypes, new_data)]
    # )

    not_changed_dtypes = [torch.uint8, torch.int8, torch.int32, torch.int64]
    data = [torch.zeros((1), dtype=dtype) for dtype in not_changed_dtypes]
    new_data = trainer.strategy.batch_to_device(data)
    assert all(val.dtype is dtype for val, dtype in zip(new_data, not_changed_dtypes)), "".join(
        [f"{dtype}: {val.dtype}" for dtype, val in zip(not_changed_dtypes, new_data)]
    )

    with pytest.raises(SystemExit):
        trainer.fit(model)

def test_device_iterations_ipu_strategy(tmpdir):
    class TestCallback(Callback):
        def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
            assert trainer.strategy.device_iterations == 2
            # assert device iterations has been set correctly within the poptorch options
            poptorch_model = trainer.strategy.poptorch_models[RunningStage.TRAINING]
            assert poptorch_model._options.toDict()["device_iterations"] == 2
            raise SystemExit

    model = IPUModel()
    trainer = Trainer(
        default_root_dir=tmpdir,
        fast_dev_run=True,
        accelerator=IPUAccelerator(),
        devices=1,
        strategy=IPUStrategy(device_iterations=2),
        callbacks=TestCallback(),
    )
    assert isinstance(trainer.strategy, IPUStrategy)
    with pytest.raises(SystemExit):
        trainer.fit(model)

def test_accumulated_batches(tmpdir):
    class TestCallback(Callback):
        def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
            # assert poptorch option have been set correctly
            poptorch_model = trainer.strategy.poptorch_models[RunningStage.TRAINING]
            assert poptorch_model._options.Training.toDict()["gradient_accumulation"] == 2
            raise SystemExit

    model = IPUModel()
    trainer = Trainer(
        default_root_dir=tmpdir,
        fast_dev_run=True,
        strategy=IPUStrategy(accelerator=IPUAccelerator()),
        devices=1,
        accumulate_grad_batches=2,
        callbacks=TestCallback(),
    )
    with pytest.raises(SystemExit):
        trainer.fit(model)

def test_stages_correct(tmpdir):
    """Ensure all stages correctly are traced correctly by asserting the output for each stage."""

    class StageModel(IPUModel):
        def training_step(self, batch, batch_idx):
            loss = super().training_step(batch, batch_idx)
            # tracing requires a loss value that depends on the model.
            # force it to be a value but ensure we use the loss.
            return (loss - loss) + torch.tensor(1)

        def validation_step(self, batch, batch_idx):
            loss = super().validation_step(batch, batch_idx)
            return (loss - loss) + torch.tensor(2)

        def test_step(self, batch, batch_idx):
            loss = super().validation_step(batch, batch_idx)
            return (loss - loss) + torch.tensor(3)

        def predict_step(self, batch, batch_idx, dataloader_idx=0):
            output = super().predict_step(batch, batch_idx)
            return (output - output) + torch.tensor(4)

    class TestCallback(Callback):
        def on_train_batch_end(self, trainer, pl_module, outputs, *_) -> None:
            assert outputs["loss"].item() == 1

        def on_validation_batch_end(self, trainer, pl_module, outputs, *_) -> None:
            assert outputs.item() == 2

        def on_test_batch_end(self, trainer, pl_module, outputs, *_) -> None:
            assert outputs.item() == 3

        def on_predict_batch_end(self, trainer, pl_module, outputs, *_) -> None:
            assert torch.all(outputs == 4).item()

    model = StageModel()
    trainer = Trainer(
        default_root_dir=tmpdir, fast_dev_run=True, accelerator=IPUAccelerator(), devices=1, callbacks=TestCallback()
    )
    trainer.fit(model)
    trainer.test(model)
    trainer.validate(model)
    trainer.predict(model, model.test_dataloader())

@pytest.mark.skip()  # todo: did not raise exception
def test_clip_gradients_fails(tmpdir):
    model = IPUModel()
    trainer = Trainer(default_root_dir=tmpdir, accelerator=IPUAccelerator(), devices=1, gradient_clip_val=10)
    with pytest.raises(MisconfigurationException, match="IPUs currently do not support clipping gradients."):
        trainer.fit(model)

@pytest.mark.xfail(RuntimeError, reason="element 0 of tensors does not require grad and does not have ...")  # todo
def test_autoreport(tmpdir):
    """Ensure autoreport dumps to a file."""
    model = IPUModel()
    autoreport_path = os.path.join(tmpdir, "report/")
    trainer = Trainer(
        default_root_dir=tmpdir,
        accelerator=IPUAccelerator(),
        devices=1,
        fast_dev_run=True,
        strategy=IPUStrategy(autoreport=True, autoreport_dir=autoreport_path),
    )
    trainer.fit(model)
    assert os.path.exists(autoreport_path)
    assert os.path.isfile(autoreport_path + "training/profile.pop")

@pytest.mark.xfail(RuntimeError, reason="element 0 of tensors does not require grad and does not have ...")  # todo
def test_manual_poptorch_dataloader(tmpdir):
    model_options = poptorch.Options()

    class IPUTestModel(IPUModel):
        def train_dataloader(self):
            dataloader = super().train_dataloader()
            # save to instance to compare the reference later
            self.poptorch_dataloader = poptorch.DataLoader(model_options, dataloader.dataset, drop_last=True)
            return self.poptorch_dataloader

    model = IPUTestModel()
    other_options = poptorch.Options()
    trainer = Trainer(
        default_root_dir=tmpdir,
        fast_dev_run=True,
        accelerator=IPUAccelerator(),
        devices=2,
        strategy=IPUStrategy(training_opts=other_options),
    )
    trainer.fit(model)

    assert isinstance(trainer.strategy, IPUStrategy)
    assert trainer.strategy.training_opts is other_options
    dataloader = trainer.train_dataloader
    assert dataloader is model.poptorch_dataloader  # exact object, was not recreated
    # dataloader uses the options in the model, not the strategy
    assert dataloader.options is model_options
    assert dataloader.options is not other_options
    assert dataloader.drop_last  # was kept

@pytest.mark.xfail(RuntimeError, reason="element 0 of tensors does not require grad and does not have ...")  # todo
def test_manual_poptorch_opts(tmpdir):
    """Ensure if the user passes manual poptorch Options, we run with the correct object."""
    model = IPUModel()
    inference_opts = poptorch.Options()
    training_opts = poptorch.Options()

    trainer = Trainer(
        default_root_dir=tmpdir,
        accelerator=IPUAccelerator(),
        devices=2,
        fast_dev_run=True,
        strategy=IPUStrategy(inference_opts=inference_opts, training_opts=training_opts),
    )
    trainer.fit(model)

    assert isinstance(trainer.strategy, IPUStrategy)
    assert trainer.strategy.training_opts == training_opts
    assert trainer.strategy.inference_opts == inference_opts

    dataloader = trainer.train_dataloader
    assert isinstance(dataloader, poptorch.DataLoader)
    assert dataloader.options == training_opts
    assert trainer.num_devices > 1  # testing this only makes sense in a distributed setting
    assert not isinstance(dataloader.sampler, DistributedSampler)

def test_manual_poptorch_opts_custom(tmpdir):
    """Ensure if the user passes manual poptorch Options with custom parameters set.

    We respect them in our poptorch options and the dataloaders.

    """
    model = IPUModel()
    training_opts = poptorch.Options()
    training_opts.deviceIterations(8)
    training_opts.replicationFactor(2)
    training_opts.Training.gradientAccumulation(2)

    inference_opts = poptorch.Options()
    inference_opts.deviceIterations(16)
    inference_opts.replicationFactor(1)
    inference_opts.Training.gradientAccumulation(1)

    class TestCallback(Callback):
        def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
            # ensure dataloaders were correctly set up during training.
            strategy = trainer.strategy
            assert isinstance(strategy, IPUStrategy)
            assert strategy.training_opts.replication_factor == 2
            assert strategy.inference_opts.replication_factor == 1

            val_dataloader = trainer.val_dataloaders
            train_dataloader = trainer.train_dataloader
            assert isinstance(val_dataloader, poptorch.DataLoader)
            assert isinstance(train_dataloader, poptorch.DataLoader)
            assert train_dataloader.options.replication_factor == 2
            assert val_dataloader.options.replication_factor == 1

    strategy = IPUStrategy(inference_opts=inference_opts, training_opts=training_opts)
    # ensure we default to the training options replication factor
    assert strategy.replication_factor == 2
    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, strategy=strategy, callbacks=TestCallback())
    trainer.fit(model)

    strategy = trainer.strategy
    assert isinstance(strategy, IPUStrategy)

    training_opts = strategy.training_opts
    assert training_opts.device_iterations == 8
    assert training_opts.replication_factor == 2
    assert training_opts.Training.gradient_accumulation == 2

    inference_opts = strategy.inference_opts
    assert inference_opts.device_iterations == 16
    assert inference_opts.replication_factor == 1
    assert inference_opts.Training.gradient_accumulation == 1

def test_replication_factor(tmpdir):
    """Ensure if the user passes manual poptorch Options with custom parameters set.

    We set them correctly in the dataloaders.

    """
    strategy = IPUStrategy()
    trainer = Trainer(
        accelerator=IPUAccelerator(), devices=2, default_root_dir=tmpdir, fast_dev_run=True, strategy=strategy
    )
    assert isinstance(trainer.accelerator, IPUAccelerator)
    assert trainer.num_devices == 2
    assert trainer.strategy.replication_factor == 2

    model = BoringModel()
    training_opts = poptorch.Options()
    inference_opts = poptorch.Options()
    training_opts.replicationFactor(8)
    inference_opts.replicationFactor(7)
    strategy = IPUStrategy(inference_opts=inference_opts, training_opts=training_opts)

    trainer = Trainer(default_root_dir=tmpdir, accelerator=IPUAccelerator(), devices=1, strategy=strategy)
    trainer.optimizers = model.configure_optimizers()[0]
    strategy._lightning_module = model
    model.trainer = trainer
    trainer.state.fn = TrainerFn.FITTING
    trainer.strategy.setup(trainer)

    trainer.state.stage = RunningStage.TRAINING
    assert trainer.strategy.replication_factor == 8
    trainer.state.stage = RunningStage.VALIDATING
    assert trainer.strategy.replication_factor == 7

    for fn, stage in (
        (TrainerFn.VALIDATING, RunningStage.VALIDATING),
        (TrainerFn.TESTING, RunningStage.TESTING),
        (TrainerFn.PREDICTING, RunningStage.PREDICTING),
    ):
        trainer.state.fn = fn
        trainer.state.stage = stage
        trainer.strategy.setup(trainer)
        assert trainer.strategy.replication_factor == 7

@pytest.mark.xfail()  # todo: assert strategy is COU
def test_default_opts(tmpdir):
    """Ensure default opts are set correctly in the IPUStrategy."""
    model = IPUModel()

    trainer = Trainer(default_root_dir=tmpdir, accelerator=IPUAccelerator(), devices=1, fast_dev_run=True)
    trainer.fit(model)
    assert isinstance(trainer.strategy, IPUStrategy)
    inference_opts = trainer.strategy.inference_opts
    training_opts = trainer.strategy.training_opts
    for opts in (inference_opts, training_opts):
        assert isinstance(opts, poptorch.Options)
        assert opts.Training.gradient_accumulation == 1
        assert opts.device_iterations == 1
        assert opts.replication_factor == 1

@pytest.mark.skip()  # todo: did not raise exception
def test_multi_optimizers_fails(tmpdir):
    """Ensure if there are multiple optimizers, we throw an exception."""

    class TestModel(IPUModel):
        def configure_optimizers(self):
            return [torch.optim.Adam(self.parameters()), torch.optim.Adam(self.parameters())]

    model = TestModel()
    # Must switch to manual optimization mode, otherwise we would get a different error
    # (multiple optimizers only supported with manual optimization)
    model.automatic_optimization = False

    trainer = Trainer(default_root_dir=tmpdir, accelerator=IPUAccelerator(), devices=1)
    with pytest.raises(MisconfigurationException, match="IPUs currently only support one optimizer."):
        trainer.fit(model)

def test_precision_plugin():
    """Ensure precision plugin value is set correctly."""
    plugin = IPUPrecision(precision="16-mixed")
    assert plugin.precision == "16-mixed"

def test_accelerator_ipu():
    trainer = Trainer(accelerator=IPUAccelerator(), devices=1)
    assert isinstance(trainer.accelerator, IPUAccelerator)

    trainer = Trainer(accelerator=IPUAccelerator())
    assert isinstance(trainer.accelerator, IPUAccelerator)

    # TODO
    # trainer = Trainer(accelerator="auto", devices=8)
    # assert isinstance(trainer.accelerator, IPUAccelerator)

def test_accelerator_ipu_with_devices():
    trainer = Trainer(strategy=IPUStrategy(accelerator=IPUAccelerator()), devices=8)
    assert isinstance(trainer.strategy, IPUStrategy)
    assert isinstance(trainer.accelerator, IPUAccelerator)
    assert trainer.num_devices == 8

@pytest.mark.xfail(AssertionError, reason="not implemented on PL side")
def test_accelerator_auto_with_devices_ipu():
    trainer = Trainer(accelerator="auto", devices=8)
    assert isinstance(trainer.accelerator, IPUAccelerator)
    assert trainer.num_devices == 8

def test_strategy_choice_ipu_strategy():
    trainer = Trainer(strategy=IPUStrategy(), accelerator=IPUAccelerator(), devices=8)
    assert isinstance(trainer.strategy, IPUStrategy)

def test_device_type_when_ipu_strategy_passed():
    trainer = Trainer(strategy=IPUStrategy(), accelerator=IPUAccelerator(), devices=8)
    assert isinstance(trainer.strategy, IPUStrategy)
    assert isinstance(trainer.accelerator, IPUAccelerator)

def test_poptorch_models_at_different_stages(tmpdir):
    strategy = IPUStrategy()
    trainer = Trainer(default_root_dir=tmpdir, strategy=strategy, accelerator=IPUAccelerator(), devices=8)
    model = BoringModel()
    model.trainer = trainer
    strategy._lightning_module = model

    trainer.optimizers = model.configure_optimizers()[0]
    trainer.state.fn = TrainerFn.FITTING
    trainer.strategy.setup(trainer)
    assert list(trainer.strategy.poptorch_models) == [
        RunningStage.TRAINING,
        RunningStage.VALIDATING,
        RunningStage.SANITY_CHECKING,
    ]

    for fn, stage in (
        (TrainerFn.VALIDATING, RunningStage.VALIDATING),
        (TrainerFn.TESTING, RunningStage.TESTING),
        (TrainerFn.PREDICTING, RunningStage.PREDICTING),
    ):
        trainer.state.fn = fn
        trainer.state.stage = stage
        trainer.strategy.setup(trainer)
        assert list(trainer.strategy.poptorch_models) == [stage]

@pytest.mark.xfail(AssertionError, reason="not implemented on PL side")
def test_devices_auto_choice_ipu():
    trainer = Trainer(accelerator="auto", devices="auto")
    assert trainer.num_devices == 4
    assert isinstance(trainer.accelerator, IPUAccelerator)

@pytest.mark.parametrize("trainer_kwargs", [{"accelerator": "ipu"}])
@pytest.mark.parametrize("hook", ["transfer_batch_to_device", "on_after_batch_transfer"])
def test_raise_exception_with_batch_transfer_hooks(monkeypatch, hook, trainer_kwargs, tmpdir):
    """Test that an exception is raised when overriding batch_transfer_hooks."""

    def custom_method(self, batch, *_, **__):
        return batch + 1

    trainer = Trainer(default_root_dir=tmpdir, accelerator="ipu")

    model = BoringModel()
    setattr(model, hook, custom_method)

    match_pattern = rf"Overriding `{hook}` is not .* with IPUs"
    with pytest.raises(MisconfigurationException, match=match_pattern):
        trainer.fit(model)

def test_unsupported_ipu_choice(monkeypatch):
    with pytest.raises(ValueError, match=r"accelerator='ipu', precision='bf16-mixed'\)` is not supported"):
        Trainer(accelerator="ipu", precision="bf16-mixed")
    with pytest.raises(ValueError, match=r"accelerator='ipu', precision='64-true'\)` is not supported"):
        Trainer(accelerator="ipu", precision="64-true")

@pytest.mark.xfail()  # todo
def test_num_stepping_batches_with_ipu(monkeypatch):
    """Test stepping batches with IPU training which acts like DP."""
    import lightning.pytorch.strategies.ipu as ipu

    monkeypatch.setattr(ipu, "_IPU_AVAILABLE", True)
    trainer = Trainer(accelerator="ipu", devices=2, max_epochs=1)
    model = BoringModel()
    trainer._data_connector.attach_data(model)
    trainer.strategy.connect(model)
    assert isinstance(trainer.strategy, IPUStrategy)
    assert trainer.estimated_stepping_batches == 64


from typing import Tuple

from django.contrib.auth import get_user_model
from django.contrib.auth.models import User
from django.core.paginator import Paginator
from django.db import models
from django.db.models import QuerySet

class Review(models.Model):
    """
    Review Model
    """

    RATING_CHOICES = zip(range(1, 6), range(1, 6))

    reviewer = models.ForeignKey(get_user_model(), on_delete=models.CASCADE)
    book = models.ForeignKey('book.Book', on_delete=models.CASCADE)

    title = models.CharField(max_length=128)
    text = models.TextField()
    rating = models.IntegerField(choices=RATING_CHOICES)

    approved = models.BooleanField(default=False)

    def __repr__(self) -> str:
        """
        Repr method overriding
        :return: str
        """
        return self.reviewer.username + ' : ' + self.book.title

    def __str__(self) -> str:
        """
        Str method overriding
        :return: str
        """
        return self.title

    @property
    def intro(self) -> str:
        """
        Review Intro displayed in Book Details View
        :return: str
        """
        return self.text[:100] + '...'

    @property
    def owner(self) -> User:
        """
        Property for accessing the user who wrote the review.
        Added for compatibility with Comment model.
        :return: User
        """
        return self.reviewer

    def save(self, *args, **kwargs):
        """
        All comments written by superuser or moderator are approved by default
        """
        if self.reviewer.is_superuser or self.reviewer.groups.filter(
                name='moderator').exists():
            self.approved = True
        super().save(*args, **kwargs)

    @classmethod
    def get_book_reviews(cls, *, pk: int, page: int = 0,
                         paginate_by: int = 12) -> Tuple[QuerySet, Paginator]:
        """
        Return Book Reviews and Pagination object.
        Used in Book Details view.
        :param pk: int
        :param page: int = 0
        :param paginate_by: int = 12
        :return: Tuple[QuerySet, Paginator]
        """
        pag = Paginator(
            cls.objects.select_related('reviewer').filter(approved=True,
                                                          book__pk=pk),
            paginate_by)
        return pag.page(page).object_list, pag

class Comment(models.Model):
    """
    Comment Model
    """

    writer = models.ForeignKey(get_user_model(), on_delete=models.CASCADE)
    review = models.ForeignKey('comment.Review', on_delete=models.CASCADE)
    replied_to = models.ForeignKey('comment.Comment', on_delete=models.CASCADE,
                                   null=True, blank=True, default=None)
    text = models.CharField(max_length=1024)
    approved = models.BooleanField(default=False)

    @property
    def owner(self) -> User:
        """
        Property for accessing the user who wrote the review.
        Added for compatibility with Review model.
        :return: User
        """
        return self.writer

    def __repr__(self) -> str:
        """
        Repr method overriding
        :return: str
        """
        return self.writer.username + " : " + str(self.review.pk)

    def __str__(self) -> str:
        """
        Str method overriding
        :return: str
        """
        return self.text[:100]

    def save(self, *args, **kwargs):
        """
        All comments written by superuser or moderator are approved by default
        """
        if self.writer.is_superuser or self.writer.groups.filter(
                name='moderator').exists():
            self.approved = True
        super().save(*args, **kwargs)

    @classmethod
    def get_review_comments(cls, *, pk: int, page: int = 0,
                            paginate_by: int = 12) \
            -> Tuple[QuerySet, Paginator]:
        """
        Return Review Comments and Pagination object.
        Used in Book Reviews view.
        :param pk: int
        :param page: int = 0
        :param paginate_by: int = 12
        :return: Tuple[QuerySet, Paginator]
        """
        pag = Paginator(cls.objects.prefetch_related('writer')
                        .filter(approved=True, review__pk=pk), paginate_by)
        return pag.page(page).object_list, pag