//! A native-gate circuit service in turns: the gates run as written.
//!
//! - **Auth**: `Authorization: apiKey <key>`.
//! - **Device**: `GET /backends/{backend}`: its `qubits`. Every pair of a
//!   trapped-ion register interacts, so the map is complete and routing never
//!   inserts a `SWAP`.
//! - **Submit**: `POST /jobs` of the circuit job type with the `native`
//!   gateset (`gpi`, `gpi2`, `zz`), phases and angles in turns, so every dyadic
//!   angle crosses the wire exactly; debiasing off, so the job runs the
//!   program as compiled. On the simulator, `noise.model` picks a device's
//!   noise.
//! - **Results**: `GET /jobs/{id}/results/probabilities`: probabilities keyed by
//!   decimal integers whose bit `p` is physical qubit `p`, read back onto
//!   logical qubits through the program's final layout, and made counts by
//!   largest remainders.
//!
//! The service reports an energy estimate per job (`stats.kwh`); it is kept as
//! the service's figure ([`JobInfo::service_kwh`]).

use crate::transport::{Request, Transport};
use crate::{Device, JobInfo, JobStatus, Provider, QpuError, Secret, counts_from_probabilities, field, field_str, json_ok, logical_index};
use serde_json::{Value, json};
use wai_quantum::quantum_fabric::Provenance;
use wai_quantum::quantum_fabric::sampling::Counts;
use wai_quantum::quantum_toolchain::native::{Basis, Coupling, Dyadic, NativeGate, NativeProgram};

/// The service's API base.
pub const BASE: &str = "https://api.ionq.co/v0.4"; // allow-name: the service's endpoint
/// The job type of a gate-list circuit.
pub const JOB_TYPE: &str = "ionq.circuit.v1"; // allow-name: a protocol job type

/// A client for one backend: `"simulator"` or a processor such as
/// a processor's backend name.
pub struct NativeTurns<T: Transport> {
    transport: T,
    key: Secret,
    backend: String,
    noise: Option<String>,
    base: String,
}

impl<T: Transport> NativeTurns<T> {
    pub fn new(transport: T, key: Secret, backend: impl Into<String>) -> Self {
        NativeTurns { transport, key, backend: backend.into(), noise: None, base: BASE.into() }
    }

    /// From `WAI_QPU_TURNS_KEY` and `WAI_QPU_TURNS_BACKEND` (default
    /// `simulator`), with `WAI_QPU_TURNS_NOISE` as the simulator's noise model
    /// when set.
    pub fn from_env(transport: T) -> Result<Self, QpuError> {
        let backend = std::env::var("WAI_QPU_TURNS_BACKEND").ok().filter(|b| !b.is_empty()).unwrap_or_else(|| "simulator".into());
        let mut c = NativeTurns::new(transport, Secret::from_env("WAI_QPU_TURNS_KEY")?, backend);
        c.noise = std::env::var("WAI_QPU_TURNS_NOISE").ok().filter(|n| !n.is_empty());
        Ok(c)
    }

    /// The simulator with a device's noise model, by the name the service
    /// gives it.
    pub fn with_noise(mut self, model: impl Into<String>) -> Self {
        self.noise = Some(model.into());
        self
    }

    /// Another endpoint: a test server.
    pub fn with_base(mut self, base: impl Into<String>) -> Self {
        self.base = base.into();
        self
    }

    fn authed(&self, r: Request) -> Request {
        r.header("Authorization", format!("apiKey {}", self.key.expose())).header("Accept", "application/json")
    }

    fn is_simulator(&self) -> bool {
        self.backend == "simulator"
    }
}

/// A dyadic angle in turns, `m/2^k`: exact in binary64.
pub fn turns(d: Dyadic) -> f64 {
    let (m, k) = d.turns();
    m as f64 / (1u64 << k) as f64
}

fn complete(n: u16) -> Coupling {
    let mut pairs = Vec::new();
    for a in 0..n {
        for b in a + 1..n {
            pairs.push((a, b));
        }
    }
    Coupling::new(n, &pairs).expect("a complete map is valid")
}

impl<T: Transport> Provider for NativeTurns<T> {
    fn device(&self) -> Result<Device, QpuError> {
        let r = self.transport.send(&self.authed(Request::get(format!("{}/backends/{}", self.base, self.backend))))?;
        let v = json_ok(&r)?;
        let n = field(&v, &["qubits"])?.as_u64().filter(|&n| n > 0 && n <= 64).ok_or_else(|| QpuError::Protocol("qubits".into()))?;
        let provenance = if self.is_simulator() {
            Provenance::Simulator(format!("native-turns simulator, noise {}", self.noise.as_deref().unwrap_or("ideal")))
        } else {
            Provenance::Silicon
        };
        Ok(Device { name: format!("native-turns:{}", self.backend), basis: Basis::GpiZz, coupling: complete(n as u16), provenance })
    }

    fn submit(&self, program: &NativeProgram, shots: u64) -> Result<String, QpuError> {
        if program.basis != Basis::GpiZz {
            return Err(QpuError::Config("this service runs gpi, gpi2 and zz".into()));
        }
        let mut circuit: Vec<Value> = Vec::with_capacity(program.gates.len());
        let mut width = program.final_layout.iter().copied().max().map_or(0, |q| q + 1);
        for g in &program.gates {
            circuit.push(match *g {
                NativeGate::Gpi(q, p) => {
                    width = width.max(q + 1);
                    json!({"gate": "gpi", "target": q, "phase": turns(p)})
                }
                NativeGate::Gpi2(q, p) => {
                    width = width.max(q + 1);
                    json!({"gate": "gpi2", "target": q, "phase": turns(p)})
                }
                NativeGate::Zz(a, b, t) => {
                    width = width.max(a.max(b) + 1);
                    json!({"gate": "zz", "targets": [a, b], "angle": turns(t)})
                }
                other => return Err(QpuError::Config(format!("{other:?} is not one of this service's native gates"))),
            });
        }
        let mut body = json!({
            "type": JOB_TYPE,
            "backend": self.backend,
            "shots": shots,
            "settings": {"error_mitigation": {"debiasing": false}},
            "input": {"gateset": "native", "qubits": width, "circuit": circuit}
        });
        if self.is_simulator()
            && let Some(model) = &self.noise
        {
            body["noise"] = json!({"model": model});
        }
        let r = self.transport.send(&self.authed(Request::post(format!("{}/jobs", self.base), body.to_string().into_bytes()).header("Content-Type", "application/json")))?;
        Ok(field_str(&json_ok(&r)?, &["id"])?.to_string())
    }

    fn status(&self, id: &str) -> Result<JobStatus, QpuError> {
        let r = self.transport.send(&self.authed(Request::get(format!("{}/jobs/{id}", self.base))))?;
        let v = json_ok(&r)?;
        Ok(match field_str(&v, &["status"])? {
            "submitted" | "ready" | "started" => JobStatus::Pending,
            "completed" => JobStatus::Done,
            other => JobStatus::Ended { status: other.to_string(), reason: v.get("failure").map(|f| f.to_string()).unwrap_or_default() },
        })
    }

    fn counts(&self, id: &str, program: &NativeProgram, shots: u64) -> Result<(Counts, JobInfo), QpuError> {
        let job = json_ok(&self.transport.send(&self.authed(Request::get(format!("{}/jobs/{id}", self.base))))?)?;
        let service_kwh = job.get("stats").and_then(|s| s.get("kwh")).and_then(|k| k.as_f64());
        let r = self.transport.send(&self.authed(Request::get(format!("{}/jobs/{id}/results/probabilities", self.base))))?;
        let v = json_ok(&r)?;
        let map = v.as_object().ok_or_else(|| QpuError::Protocol("probabilities is not a map".into()))?;
        let mut p = vec![0.0f64; 1 << program.n_qubits];
        for (key, prob) in map {
            let physical: u64 = key.parse().map_err(|_| QpuError::Counts(format!("outcome {key:?} is not a decimal integer")))?;
            let prob = prob.as_f64().filter(|x| x.is_finite() && *x >= 0.0).ok_or_else(|| QpuError::Counts(format!("probability of {key:?}")))?;
            p[logical_index(physical, &program.final_layout)] += prob;
        }
        Ok((counts_from_probabilities(&p, shots), JobInfo { service_kwh }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn turns_are_exact_binary_fractions() {
        let d = |m, k| Dyadic::new(m, k).unwrap();
        assert_eq!(turns(Dyadic::ZERO), 0.0);
        assert_eq!(turns(d(1, 2)), 0.25);
        assert_eq!(turns(d(3, 2)), 0.75);
        assert_eq!(turns(d(1, 32)), 1.0 / 4294967296.0);
        assert_eq!(json!(turns(d(5, 4))).to_string(), "0.3125");
    }

    #[test]
    fn every_pair_of_the_register_interacts() {
        let c = complete(4);
        for a in 0..4 {
            for b in 0..4 {
                assert_eq!(c.connected(a, b), a != b);
            }
        }
    }
}