// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use core::ops::Deref;
use serde::{Deserialize, Serialize};

use super::super::{
    compiled_expression::CompiledExpression,
    error::IO::{self, CannotParseJSON},
    stages::tfactory::{
        TFactoryDistillationUnitResources, TFactoryDistillationUnitTemplate,
        TFactoryDistillationUnitType, TFactoryFormula,
    },
    Error, Result,
};

#[derive(Default, Serialize, Deserialize)]
pub struct TFactoryDistillationUnitSpecifications(Vec<TFactoryDistillationUnitSpecification>);

impl TFactoryDistillationUnitSpecifications {
    pub fn as_templates(&self) -> Result<Vec<TFactoryDistillationUnitTemplate>> {
        if self.is_empty() {
            Ok(TFactoryDistillationUnitTemplate::default_distillation_unit_templates())
        } else {
            self.iter()
                .map(TFactoryDistillationUnitTemplate::try_from)
                .collect()
        }
    }
}

impl Deref for TFactoryDistillationUnitSpecifications {
    type Target = Vec<TFactoryDistillationUnitSpecification>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(untagged)]
pub enum TFactoryDistillationUnitSpecification {
    #[serde(rename_all = "camelCase")]
    Predefined { name: String },
    #[serde(rename_all = "camelCase")]
    Custom {
        /// The distillation unit output name.
        display_name: String,
        /// The number of input t states accepted by the distillation unit.
        num_input_ts: u64,
        /// The number of output t states generated by the distillation unit.
        num_output_ts: u64,
        /// The failure probability formula.
        failure_probability_formula: String,
        /// The output error rate formula.
        output_error_rate_formula: String,
        /// Specification for the physical qubit protocol.
        #[serde(skip_serializing_if = "Option::is_none")]
        physical_qubit_specification: Option<TFactoryProtocolSpecificDistillationUnitSpecification>,
        /// Specification for the logical qubit protocol.
        #[serde(skip_serializing_if = "Option::is_none")]
        logical_qubit_specification: Option<TFactoryProtocolSpecificDistillationUnitSpecification>,
        /// Specification for the logical qubit protocol if necessary to override for the first round of distillation.
        #[serde(skip_serializing_if = "Option::is_none")]
        logical_qubit_specification_first_round_override:
            Option<TFactoryProtocolSpecificDistillationUnitSpecification>,
    },
}

impl TryFrom<&TFactoryDistillationUnitSpecification> for TFactoryDistillationUnitTemplate {
    type Error = Error;

    fn try_from(specification: &TFactoryDistillationUnitSpecification) -> Result<Self> {
        match specification {
            TFactoryDistillationUnitSpecification::Predefined { name } => {
                Ok(Self::from_name(name)?)
            }
            TFactoryDistillationUnitSpecification::Custom {
                display_name,
                num_input_ts,
                num_output_ts,
                failure_probability_formula,
                output_error_rate_formula,
                physical_qubit_specification,
                logical_qubit_specification,
                logical_qubit_specification_first_round_override,
            } => {
                // Validate the specification.
                if *num_input_ts == 0 {
                    return Err(CannotParseJSON(serde::de::Error::custom(
                        "The number of input t states must be greater than 0.",
                    ))
                    .into());
                }

                if *num_output_ts == 0 {
                    return Err(CannotParseJSON(serde::de::Error::custom(
                        "The number of output t states must be greater than 0.",
                    ))
                    .into());
                }

                if physical_qubit_specification.is_none() && logical_qubit_specification.is_none() {
                    return Err(CannotParseJSON(serde::de::Error::custom(
                        "At least one of the physical or the logical specification must be specified.",
                    )).into());
                }

                let unit_type = if physical_qubit_specification.is_some()
                    && logical_qubit_specification.is_some()
                {
                    TFactoryDistillationUnitType::Combined
                } else if physical_qubit_specification.is_some() {
                    if logical_qubit_specification_first_round_override.is_some() {
                        return Err(CannotParseJSON(
                            serde::de::Error::custom(
                                "The logical specification can be overridden for the first round of distillation only if the logical specification is specified.",
                            ),
                        ).into());
                    }

                    TFactoryDistillationUnitType::Physical
                } else {
                    TFactoryDistillationUnitType::Logical
                };

                let physical_qubit_specification = if let Some(spec) = physical_qubit_specification
                {
                    Some(spec.as_resources()?)
                } else {
                    None
                };
                let logical_qubit_specification = if let Some(spec) = logical_qubit_specification {
                    Some(spec.as_resources()?)
                } else {
                    None
                };
                let logical_qubit_specification_first_round_override =
                    if let Some(spec) = logical_qubit_specification_first_round_override {
                        Some(spec.as_resources()?)
                    } else {
                        None
                    };

                let failure_probability_function =
                    TFactoryFormula::from(CompiledExpression::from_string(
                        failure_probability_formula,
                        "failure_probability_formula",
                    )?);

                let output_error_rate_function =
                    TFactoryFormula::from(CompiledExpression::from_string(
                        output_error_rate_formula,
                        "output_error_rate_formula",
                    )?);

                Ok(Self {
                    name: display_name.clone(),
                    num_input_ts: *num_input_ts,
                    num_output_ts: *num_output_ts,
                    failure_probability_function,
                    output_error_rate_function,
                    unit_type,
                    physical_qubit_specification,
                    logical_qubit_specification,
                    logical_qubit_specification_first_round_override,
                })
            }
        }
    }
}

#[derive(Default, Debug, Serialize, Deserialize, Clone)]
#[serde(rename_all(serialize = "camelCase", deserialize = "camelCase"))]
pub struct TFactoryProtocolSpecificDistillationUnitSpecification {
    /// The number of unit qubits utilized in distillation.
    pub(crate) num_unit_qubits: u64,
    /// The duration of distillation measured in qubit cycles.
    pub(crate) duration_in_qubit_cycle_time: u64,
}

impl TFactoryProtocolSpecificDistillationUnitSpecification {
    pub fn as_resources(&self) -> core::result::Result<TFactoryDistillationUnitResources, IO> {
        if self.num_unit_qubits == 0 {
            return Err(CannotParseJSON(serde::de::Error::custom(
                "The number of unit qubits must be greater than 0.",
            )));
        }

        Ok(TFactoryDistillationUnitResources {
            duration_in_qubit_cycle_time: self.duration_in_qubit_cycle_time,
            num_unit_qubits: self.num_unit_qubits,
        })
    }
}
