# Copyright 2024 - 2025 Aegiq Ltd.
#
# 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.

from copy import deepcopy
from random import random

import pytest

from lightworks import (
    Analyzer,
    Batch,
    Parameter,
    PhotonicCircuit,
    PostSelection,
    Sampler,
    Simulator,
    State,
    Unitary,
    random_unitary,
)
from lightworks.emulator import Backend


class TestBatchTask:
    """
    Unit tests for the Batch object.
    """

    def setup_method(self):
        """
        Creates components required for testing
        """
        self.circuits = [Unitary(random_unitary(8)) for _i in range(4)]
        self.inputs = [
            State([1, 1, 0, 0, 1, 1, 0, 0]),
            State([1, 1, 0, 0, 0, 0, 1, 1]),
            State([0, 0, 1, 1, 1, 1, 0, 0]),
            State([0, 0, 1, 1, 0, 0, 1, 1]),
        ]
        self.samples = [1, 100, 10000, 1000000]
        self.min_detections = [1, 2, 3, 4]
        self.post_selection = [PostSelection()] * 4

    def test_task_setup(self):
        """
        Checks basic Batch setup works correctly.
        """
        Batch(
            Sampler,
            task_args=[self.circuits, self.inputs, self.samples],
            task_kwargs={
                "min_detection": self.min_detections,
                "post_selection": self.post_selection,
            },
        )

    def test_setup_no_task(self):
        """
        Checks Batch can be initialized with no arguments.
        """
        Batch()

    def test_setup_task_initialised(self):
        """
        Checks an error is raised when an already initialized task is provided
        to Batch.
        """
        sim = Simulator(PhotonicCircuit(2), State([1, 0]))
        with pytest.raises(TypeError):
            Batch(sim)

    def test_task_setup_all_kwargs(self):
        """
        Checks basic Batch setup works correctly when all arguments are
        specified as keyword arguments.
        """
        Batch(
            Sampler,
            task_kwargs={
                "circuit": self.circuits,
                "input_state": self.inputs,
                "n_samples": self.samples,
                "min_detection": self.min_detections,
                "post_selection": self.post_selection,
            },
        )

    def test_task_setup_single_value(self):
        """
        Checks basic Batch setup works correctly when only a single value
        is provided for some of the args.
        """
        Batch(
            Sampler,
            task_args=[[self.circuits[1]], [self.inputs[2]], self.samples],
            task_kwargs={
                "min_detection": [self.min_detections[3]],
                "post_selection": [self.post_selection[0]],
            },
        )

    def test_task_number(self):
        """
        Checks that the correct number of tasks is generated by the
        Batch.
        """
        bc = Batch(
            Sampler,
            task_args=[self.circuits, self.inputs, self.samples],
            task_kwargs={"min_detection": self.min_detections},
        )
        assert bc.num == 4

    def test_add(self):
        """
        Checks add function allows a task to be added, and raised an error for
        not-tasks.
        """
        sim = Simulator(PhotonicCircuit(2), State([1, 0]))
        batch = Batch()
        batch.add(sim)
        assert batch.num == 1

    def test_add_not_task(self):
        """
        Checks an error is raised when a non-task is added to Batch.
        """
        batch = Batch()
        with pytest.raises(TypeError):
            batch.add(PhotonicCircuit(4))

    def test_task_number_all_kwargs(self):
        """
        Checks that the correct number of tasks is generated by the
        Batch when all args are specified by keyword arguments.
        """
        bc = Batch(
            Sampler,
            task_kwargs={
                "circuit": self.circuits,
                "input_state": self.inputs,
                "n_samples": self.samples,
                "min_detection": self.min_detections,
                "post_selection": self.post_selection,
            },
        )
        assert bc.num == 4

    def test_task_values(self):
        """
        Checks correct values are set in each of the tasks generated.
        """
        bc = Batch(
            Sampler,
            task_args=[self.circuits, self.inputs, self.samples],
            task_kwargs={
                "min_detection": self.min_detections,
                "post_selection": self.post_selection,
            },
        )
        for i, task in enumerate(bc.tasks):
            assert (task.circuit.U_full == self.circuits[i].U_full).all()
            assert task.input_state == self.inputs[i]
            assert task.n_samples == self.samples[i]
            assert task.min_detection == self.min_detections[i]
            assert task.post_selection == self.post_selection[i]

    def test_task_values_all_kwargs(self):
        """
        Checks correct values are set in each of the tasks generated when all
        args are provided as keyword arguments.
        """
        bc = Batch(
            Sampler,
            task_kwargs={
                "circuit": self.circuits,
                "input_state": self.inputs,
                "n_samples": self.samples,
                "min_detection": self.min_detections,
                "post_selection": self.post_selection,
            },
        )
        for i, task in enumerate(bc.tasks):
            assert (task.circuit.U_full == self.circuits[i].U_full).all()
            assert task.input_state == self.inputs[i]
            assert task.n_samples == self.samples[i]
            assert task.min_detection == self.min_detections[i]
            assert task.post_selection == self.post_selection[i]

    def test_task_with_parameter_setup(self):
        """
        Checks that parameters can be provided as part of the setup of the
        Batch.
        """
        parameters = {Parameter(1): [1, 2, 3, 4], Parameter(1): [5, 6, 7, 8]}
        Batch(
            Sampler,
            task_kwargs={
                "circuit": self.circuits,
                "input_state": self.inputs,
                "n_samples": self.samples,
                "min_detection": self.min_detections,
                "post_selection": self.post_selection,
            },
            parameters=parameters,
        )

    def test_task_creates_different_circuits(self):
        """
        Checks that parameters can be provided as part of the setup of the
        Batch and that this creates unique circuits.
        """
        # Generate parametrised circuit
        parameters = {}
        circuit = PhotonicCircuit(8)
        for i, m in enumerate([0, 2, 4, 6, 1, 3, 5, 0, 2, 4, 6]):
            theta = Parameter(1, label=f"theta_{i}")
            phi = Parameter(1, label=f"phi_{i}")
            circuit.ps(m, phi)
            circuit.bs(m)
            circuit.ps(m + 1, theta)
            circuit.bs(m)
            parameters[theta] = [6.3 * random() for i in range(4)]
            parameters[phi] = [6.3 * random() for i in range(4)]
        # Then setup batch task
        bc = Batch(
            Sampler,
            task_kwargs={
                "circuit": [circuit],
                "input_state": self.inputs,
                "n_samples": self.samples,
                "min_detection": self.min_detections,
                "post_selection": self.post_selection,
            },
            parameters=parameters,
        )
        # deepcopy tasks to ensure no modification by parameters
        all_tasks = deepcopy(bc.tasks)
        # Check all unitaries are equivalent to the expected from the circuit
        for i, task in enumerate(all_tasks):
            for p in parameters:
                p.set(parameters[p][i])
            assert (task.circuit.U.round() == circuit.U.round()).all()

    @pytest.mark.parametrize(
        "outputs",
        [
            None,
            [None],
            [State([0, 1, 0, 1])],
            [State([0, 1, 0, 1])] * 5,
            [[State([0, 1, 0, 1]), State([0, 1, 1, 0])]] * 5,
        ],
    )
    def test_simulator_batch(self, outputs):
        """
        Checks the simulator can be run as a batch.
        """
        circuits = [Unitary(random_unitary(4)) for _ in range(5)]
        inputs = [State([1, 0, 1, 0])]
        kwargs = {"outputs": outputs} if outputs is not None else {}
        batch = Batch(Simulator, [circuits, inputs], kwargs)
        Backend("permanent").run(batch)

    @pytest.mark.parametrize(
        "expected",
        [None, [None], [{State([1, 0, 1, 0]): State([1, 0, 1, 0])}]],
    )
    def test_analyzer_batch(self, expected):
        """
        Checks the simulator can be run as a batch.
        """
        circuits = [Unitary(random_unitary(4)) for _ in range(5)]
        inputs = [State([1, 0, 1, 0])]
        kwargs = {"expected": expected} if expected is not None else {}
        batch = Batch(Analyzer, [circuits, inputs], kwargs)
        Backend("permanent").run(batch)
