# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# adapted from https://github.com/learnables/learn2learn/blob/master/examples/vision/protonet_miniimagenet.py#L154

"""## Train file https://www.dropbox.com/s/9g8c6w345s2ek03/mini-imagenet-cache-train.pkl?dl=1

## Validation File
https://www.dropbox.com/s/ip1b7se3gij3r1b/mini-imagenet-cache-validation.pkl?dl=1

Followed by renaming the pickle files
cp './mini-imagenet-cache-train.pkl?dl=1' './mini-imagenet-cache-train.pkl'
cp './mini-imagenet-cache-validation.pkl?dl=1' './mini-imagenet-cache-validation.pkl'
"""

import warnings
from dataclasses import dataclass
from typing import Tuple, Union

import kornia.augmentation as Ka
import kornia.geometry as Kg
import learn2learn as l2l
import torch
import torchvision.transforms as T

import flash
from flash.core.data.io.input import DataKeys
from flash.core.data.io.input_transform import InputTransform
from flash.core.data.transforms import ApplyToKeys, kornia_collate
from flash.image import ImageClassificationData, ImageClassifier

warnings.simplefilter("ignore")

# download MiniImagenet
train_dataset = l2l.vision.datasets.MiniImagenet(root="./", mode="train", download=False)
val_dataset = l2l.vision.datasets.MiniImagenet(root="./", mode="validation", download=False)


@dataclass
class ImageClassificationInputTransform(InputTransform):

    image_size: Tuple[int, int] = (196, 196)
    mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406)
    std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225)

    def per_sample_transform(self):
        return T.Compose(
            [
                ApplyToKeys(
                    DataKeys.INPUT,
                    T.Compose(
                        [
                            T.ToTensor(),
                            Kg.Resize((196, 196)),
                            # SPATIAL
                            Ka.RandomHorizontalFlip(p=0.25),
                            Ka.RandomRotation(degrees=90.0, p=0.25),
                            Ka.RandomAffine(degrees=1 * 5.0, shear=1 / 5, translate=1 / 20, p=0.25),
                            Ka.RandomPerspective(distortion_scale=1 / 25, p=0.25),
                            # PIXEL-LEVEL
                            Ka.ColorJitter(brightness=1 / 30, p=0.25),  # brightness
                            Ka.ColorJitter(saturation=1 / 30, p=0.25),  # saturation
                            Ka.ColorJitter(contrast=1 / 30, p=0.25),  # contrast
                            Ka.ColorJitter(hue=1 / 30, p=0.25),  # hue
                            Ka.RandomMotionBlur(kernel_size=2 * (4 // 3) + 1, angle=1, direction=1.0, p=0.25),
                            Ka.RandomErasing(scale=(1 / 100, 1 / 50), ratio=(1 / 20, 1), p=0.25),
                        ]
                    ),
                ),
                ApplyToKeys(DataKeys.TARGET, torch.as_tensor),
            ]
        )

    def train_per_sample_transform(self):
        return T.Compose(
            [
                ApplyToKeys(
                    DataKeys.INPUT,
                    T.Compose(
                        [
                            T.ToTensor(),
                            T.Resize(self.image_size),
                            T.Normalize(self.mean, self.std),
                            T.RandomHorizontalFlip(),
                            T.ColorJitter(),
                            T.RandomAutocontrast(),
                            T.RandomPerspective(),
                        ]
                    ),
                ),
                ApplyToKeys("target", torch.as_tensor),
            ]
        )

    def per_batch_transform_on_device(self):
        return ApplyToKeys(
            DataKeys.INPUT,
            Ka.RandomHorizontalFlip(p=0.25),
        )

    def collate(self):
        return kornia_collate


# construct datamodule

datamodule = ImageClassificationData.from_tensors(
    train_data=train_dataset.x,
    train_targets=torch.from_numpy(train_dataset.y.astype(int)),
    val_data=val_dataset.x,
    val_targets=torch.from_numpy(val_dataset.y.astype(int)),
    train_transform=ImageClassificationInputTransform,
    val_transform=ImageClassificationInputTransform,
    batch_size=1,
)

model = ImageClassifier(
    backbone="resnet18",
    training_strategy="prototypicalnetworks",
    training_strategy_kwargs={
        "epoch_length": 10 * 16,
        "meta_batch_size": 1,
        "num_tasks": 200,
        "test_num_tasks": 2000,
        "ways": datamodule.num_classes,
        "shots": 1,
        "test_ways": 5,
        "test_shots": 1,
        "test_queries": 15,
    },
    optimizer=torch.optim.Adam,
    learning_rate=0.001,
)

trainer = flash.Trainer(
    max_epochs=1,
    gpus=1,
    precision=16,
)

trainer.finetune(model, datamodule=datamodule, strategy="no_freeze")

flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py
# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# adapted from https://github.com/learnables/learn2learn/blob/master/examples/vision/protonet_miniimagenet.py#L154

"""## Train file https://www.dropbox.com/s/9g8c6w345s2ek03/mini-imagenet-cache-train.pkl?dl=1

## Validation File
https://www.dropbox.com/s/ip1b7se3gij3r1b/mini-imagenet-cache-validation.pkl?dl=1

Followed by renaming the pickle files
cp './mini-imagenet-cache-train.pkl?dl=1' './mini-imagenet-cache-train.pkl'
cp './mini-imagenet-cache-validation.pkl?dl=1' './mini-imagenet-cache-validation.pkl'
"""

import warnings
from dataclasses import dataclass
from typing import Tuple, Union

import kornia.augmentation as Ka
import kornia.geometry as Kg
import learn2learn as l2l
import torch
import torchvision.transforms as T

import flash
from flash.core.data.io.input import DataKeys
from flash.core.data.io.input_transform import InputTransform
from flash.core.data.transforms import ApplyToKeys, kornia_collate
from flash.image import ImageClassificationData, ImageClassifier

warnings.simplefilter("ignore")

# download MiniImagenet
train_dataset = l2l.vision.datasets.MiniImagenet(root="./", mode="train", download=False)
val_dataset = l2l.vision.datasets.MiniImagenet(root="./", mode="validation", download=False)


@dataclass
class ImageClassificationInputTransform(InputTransform):

    image_size: Tuple[int, int] = (196, 196)
    mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406)
    std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225)

    def per_sample_transform(self):
        return T.Compose(
            [
                ApplyToKeys(
                    DataKeys.INPUT,
                    T.Compose(
                        [
                            T.ToTensor(),
                            Kg.Resize((196, 196)),
                            # SPATIAL
                            Ka.RandomHorizontalFlip(p=0.25),
                            Ka.RandomRotation(degrees=90.0, p=0.25),
                            Ka.RandomAffine(degrees=1 * 5.0, shear=1 / 5, translate=1 / 20, p=0.25),
                            Ka.RandomPerspective(distortion_scale=1 / 25, p=0.25),
                            # PIXEL-LEVEL
                            Ka.ColorJitter(brightness=1 / 30, p=0.25),  # brightness
                            Ka.ColorJitter(saturation=1 / 30, p=0.25),  # saturation
                            Ka.ColorJitter(contrast=1 / 30, p=0.25),  # contrast
                            Ka.ColorJitter(hue=1 / 30, p=0.25),  # hue
                            Ka.RandomMotionBlur(kernel_size=2 * (4 // 3) + 1, angle=1, direction=1.0, p=0.25),
                            Ka.RandomErasing(scale=(1 / 100, 1 / 50), ratio=(1 / 20, 1), p=0.25),
                        ]
                    ),
                ),
                ApplyToKeys(DataKeys.TARGET, torch.as_tensor),
            ]
        )

    def train_per_sample_transform(self):
        return T.Compose(
            [
                ApplyToKeys(
                    DataKeys.INPUT,
                    T.Compose(
                        [
                            T.ToTensor(),
                            T.Resize(self.image_size),
                            T.Normalize(self.mean, self.std),
                            T.RandomHorizontalFlip(),
                            T.ColorJitter(),
                            T.RandomAutocontrast(),
                            T.RandomPerspective(),
                        ]
                    ),
                ),
                ApplyToKeys("target", torch.as_tensor),
            ]
        )

    def per_batch_transform_on_device(self):
        return ApplyToKeys(
            DataKeys.INPUT,
            Ka.RandomHorizontalFlip(p=0.25),
        )

    def collate(self):
        return kornia_collate


# construct datamodule

datamodule = ImageClassificationData.from_tensors(
    train_data=train_dataset.x,
    train_targets=torch.from_numpy(train_dataset.y.astype(int)),
    val_data=val_dataset.x,
    val_targets=torch.from_numpy(val_dataset.y.astype(int)),
    train_transform=ImageClassificationInputTransform,
    val_transform=ImageClassificationInputTransform,
    batch_size=1,
)

model = ImageClassifier(
    backbone="resnet18",
    training_strategy="prototypicalnetworks",
    training_strategy_kwargs={
        "epoch_length": 10 * 16,
        "meta_batch_size": 1,
        "num_tasks": 200,
        "test_num_tasks": 2000,
        "ways": datamodule.num_classes,
        "shots": 1,
        "test_ways": 5,
        "test_shots": 1,
        "test_queries": 15,
    },
    optimizer=torch.optim.Adam,
    learning_rate=0.001,
)

trainer = flash.Trainer(
    max_epochs=1,
    gpus=1,
    accelerator="gpu",
    precision=16,
)

trainer.finetune(model, datamodule=datamodule, strategy="no_freeze")

flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py
diff --git a/README.md b/README.md
index 1734992079..26773424f4 100644
--- a/README.md
+++ b/README.md
@@ -129,7 +129,7 @@ model.serve()
 or make predictions from raw data directly.
 
 ```py
-trainer = Trainer(accelerator='ddp', gpus=2)
+trainer = Trainer(strategy='ddp', accelerator="gpu", gpus=2)
 dm = SemanticSegmentationData.from_folders(predict_folder="data/CameraRGB")
 predictions = trainer.predict(model, dm)
 ```
diff --git a/flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py b/flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py
index 8af2ca255d..b506dff891 100644
--- a/flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py
+++ b/flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py
@@ -146,6 +146,7 @@ def collate(self):
 trainer = flash.Trainer(
     max_epochs=1,
     gpus=1,
+    accelerator="gpu",
     precision=16,
 )
 
diff --git a/tests/core/utilities/test_lightning_cli.py b/tests/core/utilities/test_lightning_cli.py
index 350f1509e9..d1d96aba8c 100644
--- a/tests/core/utilities/test_lightning_cli.py
+++ b/tests/core/utilities/test_lightning_cli.py
@@ -590,8 +590,8 @@ def on_exception(self, execption):
 @pytest.mark.parametrize(
     "trainer_kwargs",
     (
-        dict(accelerator="ddp_cpu"),
-        dict(accelerator="ddp_cpu", plugins="ddp_find_unused_parameters_false"),
+        dict(accelerator="cpu", strategy="ddp"),
+        dict(accelerator="cpu", strategy="ddp", plugins="ddp_find_unused_parameters_false"),
     ),
 )
 @pytest.mark.skipif(not _PL_GREATER_EQUAL_1_4_0, reason="Bugs in PL < 1.4.0")
diff --git a/tests/image/classification/test_training_strategies.py b/tests/image/classification/test_training_strategies.py
index 3fedd7be68..82832b7e40 100644
--- a/tests/image/classification/test_training_strategies.py
+++ b/tests/image/classification/test_training_strategies.py
@@ -20,7 +20,7 @@
 
 from flash import Trainer
 from flash.core.data.io.input import DataKeys
-from flash.core.utilities.imports import _IMAGE_TESTING, _LEARN2LEARN_AVAILABLE
+from flash.core.utilities.imports import _IMAGE_TESTING, _LEARN2LEARN_AVAILABLE, _PL_GREATER_EQUAL_1_6_0
 from flash.image import ImageClassificationData, ImageClassifier
 from flash.image.classification.adapters import TRAINING_STRATEGIES
 from tests.image.classification.test_data import _rand_image
@@ -54,7 +54,7 @@ def test_learn2learn_training_strategies_registry():
     assert TRAINING_STRATEGIES.available_keys() == ["anil", "default", "maml", "metaoptnet", "prototypicalnetworks"]
 
 
-def _test_learn2learning_training_strategies(gpus, accelerator, training_strategy, tmpdir):
+def _test_learn2learning_training_strategies(gpus, training_strategy, tmpdir, accelerator=None, strategy=None):
     train_dir = Path(tmpdir / "train")
     train_dir.mkdir()
 
@@ -87,7 +87,11 @@ def _test_learn2learning_training_strategies(gpus, accelerator, training_strateg
         training_strategy_kwargs={"ways": dm.num_classes, "shots": 4, "meta_batch_size": 4},
     )
 
-    trainer = Trainer(fast_dev_run=2, gpus=gpus, accelerator=accelerator)
+    if _PL_GREATER_EQUAL_1_6_0:
+        trainer = Trainer(fast_dev_run=2, gpus=gpus, strategy=strategy)
+    else:
+        trainer = Trainer(fast_dev_run=2, gpus=gpus, accelerator=accelerator)
+
     trainer.fit(model, datamodule=dm)
 
 
@@ -95,7 +99,7 @@ def _test_learn2learning_training_strategies(gpus, accelerator, training_strateg
 @pytest.mark.parametrize("training_strategy", ["anil", "maml", "prototypicalnetworks"])
 @pytest.mark.skipif(not _LEARN2LEARN_AVAILABLE, reason="image and learn2learn libraries aren't installed.")
 def test_learn2learn_training_strategies(training_strategy, tmpdir):
-    _test_learn2learning_training_strategies(0, None, training_strategy, tmpdir)
+    _test_learn2learning_training_strategies(0, training_strategy, tmpdir, accelerator=None)
 
 
 @pytest.mark.skipif(not _LEARN2LEARN_AVAILABLE, reason="image and learn2learn libraries aren't installed.")
@@ -111,4 +115,7 @@ def test_wrongly_specified_training_strategies():
 @pytest.mark.skipif(not os.getenv("FLASH_RUNNING_SPECIAL_TESTS", "0") == "1", reason="Should run with special test")
 @pytest.mark.skipif(not _LEARN2LEARN_AVAILABLE, reason="image and learn2learn libraries aren't installed.")
 def test_learn2learn_training_strategies_ddp(tmpdir):
-    _test_learn2learning_training_strategies(2, "ddp", "prototypicalnetworks", tmpdir)
+    if _PL_GREATER_EQUAL_1_6_0:
+        _test_learn2learning_training_strategies(2, "prototypicalnetworks", tmpdir, strategy="ddp")
+    else:
+        _test_learn2learning_training_strategies(2, "prototypicalnetworks", tmpdir, accelerator="ddp")

Fix Flash CI (special examples failing)

