# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#
# 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 gc
import sys
from pathlib import Path

# Append bert directory
bert_root_path = str(Path(__file__).parent.parent)
sys.path.append(bert_root_path)
# Append tests directory
tests_root_path = str(Path(__file__).parent)
sys.path.append(tests_root_path)
import subprocess
import pytest
import numpy as np
import time
import transformers
from args import parse_bert_args
from tests.bert_tiny_test import parse_result_for_loss_accuracy, loss_going_down, accuracy_going_up
from ipu_options import get_options
from pretraining_data import get_dataloader, TFRecordPretrainingDataset


bert_root_dir = Path(__file__).parent.parent.resolve()


def run_poprun_cmdline(poprun_args, cmdline_args, script):
    cmdline_args["--wandb"] = "false"
    cmd = ["poprun"]
    cmd.extend([str(item) for sublist in poprun_args.items() for item in sublist if item != ""])
    cmd.append("python3")
    cmd.append(script)
    cmd.extend([str(item) for sublist in cmdline_args.items() for item in sublist if item != ""])
    try:
        out = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=bert_root_dir)
    except subprocess.CalledProcessError as e:
        print(f"TEST FAILED")
        print(f"stdout={e.stdout.decode('utf-8',errors='ignore')}")
        print(f"stderr={e.stderr.decode('utf-8',errors='ignore')}")
        raise
    return out, out.stdout.decode("utf-8"), out.stderr.decode("utf-8")


def dataset():
    """
    Check if the data in two instances is different
    """
    args = "--config demo_tiny_128".split()
    config = transformers.BertConfig(**(vars(parse_bert_args(args))))
    opts = get_options(config)
    loader = TFRecordPretrainingDataset(config.input_files)
    loader = get_dataloader(config, opts)

    # Save part of the data as list
    loader_list = list(loader)[0][0][0].numpy()

    # MPI to broadcast data in root=1 to root=0
    from mpi4py import MPI

    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    loader_list_copy = np.copy(loader_list)
    comm.Bcast(loader_list, root=1)

    # Assert if data broadcast to root=0 is different
    if comm.Get_rank() == 0 and not np.all(loader_list_copy == loader_list):
        print("Passed test: instances have different data")

    # Wait until both roots are finished
    time.sleep(2)


@pytest.mark.ipus(8)
def test_poprun_dataset():
    """
    Launch poprun as subprocess and assert output from poprun_utils.py
    """
    gc.collect()
    out, stdout, stderr = run_poprun_cmdline(
        {"--num-instances": 2, "--num-replicas": 2, "--ipus-per-replica": 4},
        {},
        os.path.dirname(os.path.abspath(__file__)) + "/poprun_test.py",
    )
    assert "Passed test: instances have different data" in stdout


@pytest.mark.ipus(4)
def test_poprun_loss_down_accuracy_up():
    """
    Test that for a 3 layer toy model the loss is trending downwards and the
    accuracy is trending upwards.
    """
    gc.collect()
    out, stdout, stderr = run_poprun_cmdline(
        {
            "--num-instances": 2,
            "--num-replicas": 2,
            "--ipus-per-replica": 2,
        },
        {
            "--config": "demo_tiny_128",
            "--training-steps": 25,
            "--ipus-per-replica": 2,
            "--layers-per-ipu": 0,
            "": 3,
            "--disable-progress-bar": "True",
        },
        "run_pretraining.py",
    )
    print(f"'\nOutput was:\n{stderr}")
    assert out.returncode == 0
    losses, accs = parse_result_for_loss_accuracy(stderr)
    loss_going_down(losses)
    accuracy_going_up(accs)


if __name__ == "__main__":
    dataset()