#!/usr/bin/env python3

# coding=utf-8
# Copyright 2023-present the HuggingFace Inc. 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.

import os
import tempfile
import unittest

import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

from peft import PeftModel, PolyConfig, TaskType, get_peft_model


class TestPoly(unittest.TestCase):
    def test_poly(self):
        torch.manual_seed(0)
        model_name_or_path = "google/flan-t5-small"

        atol, rtol = 1e-6, 1e-6
        r = 8  # rank of lora in poly
        n_tasks = 3  # number of tasks
        n_skills = 2  # number of skills (loras)
        n_splits = 4  # number of heads
        lr = 1e-2
        num_epochs = 10

        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        base_model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)

        peft_config = PolyConfig(
            task_type=TaskType.SEQ_2_SEQ_LM,
            poly_type="poly",
            r=r,
            n_tasks=n_tasks,
            n_skills=n_skills,
            n_splits=n_splits,
        )

        model = get_peft_model(base_model, peft_config)

        # generate some dummy data
        text = os.__doc__.splitlines()
        assert len(text) > 10
        inputs = tokenizer(text, return_tensors="pt", padding=True)
        inputs["task_ids"] = torch.arange(len(text)) % n_tasks
        inputs["labels"] = tokenizer((["A", "B"] * 100)[: len(text)], return_tensors="pt")["input_ids"]

        # simple training loop
        model.train()
        optimizer = torch.optim.Adam(model.parameters(), lr=lr)
        losses = []
        for _ in range(num_epochs):
            outputs = model(**inputs)
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            losses.append(loss.item())

        # loss improved by at least 50%
        assert losses[-1] < (0.5 * losses[0])

        # check that saving and loading works
        torch.manual_seed(0)
        model.eval()
        logits_before = model(**inputs).logits
        tokens_before = model.generate(**inputs)

        with model.disable_adapter():
            logits_disabled = model(**inputs).logits
            tokens_disabled = model.generate(**inputs)

        assert not torch.allclose(logits_before, logits_disabled, atol=atol, rtol=rtol)
        assert not torch.allclose(tokens_before, tokens_disabled, atol=atol, rtol=rtol)

        # saving and loading
        with tempfile.TemporaryDirectory() as tmp_dir:
            model.save_pretrained(tmp_dir)
            base_model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
            loaded = PeftModel.from_pretrained(base_model, tmp_dir)

        torch.manual_seed(0)
        output_after = loaded(**inputs).logits
        tokens_after = loaded.generate(**inputs)
        assert torch.allclose(logits_before, output_after, atol=atol, rtol=rtol)
        assert torch.allclose(tokens_before, tokens_after, atol=atol, rtol=rtol)