tokitai-operator 0.1.0

use tokitai_operator::backend::TensorStore;
use tokitai_operator::backend::hardware::DeviceKind;
use tokitai_operator::domain::{Domain, DomainId};
use tokitai_operator::facade::{AuditExportFormat, BackendPolicy, FacadeGraphBuilder, Tokitai};
use tokitai_operator::object::sheaf::{Cover, FiniteSite, Inclusion, OpenId, SectionTable};
use tokitai_operator::object::{ObjectMeta, Representation, Shape, Tensor};
use tokitai_operator::op::AddOp;

fn open(id: &str) -> OpenId {
    OpenId(id.to_string())
}

#[test]
fn facade_builds_plans_executes_and_audits_i64_graphs() {
    let api = Tokitai::new();
    let mut builder = FacadeGraphBuilder::new();
    let domain = DomainId::new("integer");
    let meta = ObjectMeta::tensor(
        domain.clone(),
        Shape::from(vec![4]),
        Representation::dense_cpu(),
    );
    let lhs = builder.input(meta.clone());
    let rhs = builder.input(meta);
    let add_out = builder.add_op(AddOp, &[lhs, rhs]).unwrap()[0];
    let graph = builder.build();

    let plan = api.plan(&graph).unwrap();
    let audit = api.audit(&plan);
    assert!(audit.contains("backend: cpu_scalar"));

    let mut inputs = TensorStore::new();
    inputs.insert(
        lhs,
        Tensor::dense_cpu(domain.clone(), Shape::from(vec![4]), vec![1, 2, 3, 4]),
    );
    inputs.insert(
        rhs,
        Tensor::dense_cpu(domain, Shape::from(vec![4]), vec![10, 20, 30, 40]),
    );

    let (store, report) = api.execute_i64(&graph, &plan, inputs, &[add_out]).unwrap();
    assert_eq!(store.get(add_out).unwrap().data, vec![11, 22, 33, 44]);
    assert_eq!(report.checked_outputs, vec![add_out]);
}

#[test]
fn p190_public_api_exposes_cpu_padic_constraints_and_audit_export() {
    let api = Tokitai::cpu_only();
    let q5 = api.padic_domain(5, 3).unwrap();
    let public_plan = api.plan_padic_sum_products(&q5, 0, 1, 2).unwrap();

    assert_eq!(public_plan.backend_policy, BackendPolicy::CpuOnly);
    assert_eq!(public_plan.requested_target.kind, DeviceKind::Cpu);
    assert_eq!(public_plan.selected_target.kind, DeviceKind::Cpu);
    assert!(public_plan.fallback_reason.is_none());
    assert!(public_plan.conformance.is_none());
    assert_eq!(
        public_plan
            .plan
            .resources
            .padic
            .as_ref()
            .unwrap()
            .backend_capability,
        "padic:fixed_precision"
    );
    assert!(
        api.plan_fingerprint(&public_plan.plan)
            .starts_with("sha256:")
    );

    let lhs = vec![q5.element(1), q5.element(25), q5.element(5), q5.element(2)];
    let rhs = vec![q5.element(2), q5.element(5), q5.element(25), q5.element(3)];
    let mut store = TensorStore::new();
    store.insert(0, Tensor::dense_cpu(q5.id(), Shape::from(vec![4]), lhs));
    store.insert(1, Tensor::dense_cpu(q5.id(), Shape::from(vec![4]), rhs));

    let report = api
        .execute_padic_sum_products(&public_plan, &mut store)
        .unwrap();
    assert_eq!(report.skipped_terms, 2);
    assert!(store.contains(2));

    let audit = api.export_audit(&public_plan.plan, AuditExportFormat::Json);
    assert!(audit.contains("\"math_bound\""));
    assert!(audit.contains("\"p_adic_constraint_fingerprint\""));
}

#[test]
fn p190_public_api_prefers_gpu_but_falls_back_with_conformance_evidence() {
    let api = Tokitai::prefer_gpu_with_cpu_fallback();
    let q5 = api.padic_domain(5, 3).unwrap();
    let public_plan = api.plan_padic_sum_products(&q5, 0, 1, 2).unwrap();

    assert_eq!(
        public_plan.backend_policy,
        BackendPolicy::PreferGpuWithCpuFallback
    );
    assert_eq!(public_plan.requested_target.kind, DeviceKind::Gpu);
    assert_eq!(public_plan.selected_target.kind, DeviceKind::Cpu);
    assert!(
        public_plan
            .fallback_reason
            .as_ref()
            .unwrap()
            .contains("gpu_scaffold does not support fixed-precision p-adic execution")
    );
    let conformance = public_plan.conformance.as_ref().unwrap();
    assert!(conformance.passed());
    assert_eq!(conformance.fallback_cases().len(), 6);
}

#[test]
fn p251_public_prefer_gpu_dense_integer_add_falls_back_to_cpu_without_real_backend_feature() {
    let api = Tokitai::prefer_gpu_with_cpu_fallback();
    let mut builder = FacadeGraphBuilder::new();
    let meta = ObjectMeta::tensor(
        DomainId::new("integer"),
        Shape::from(vec![4]),
        Representation::dense_cpu(),
    );
    let lhs = builder.input(meta.clone());
    let rhs = builder.input(meta);
    let out = builder.add_op(AddOp, &[lhs, rhs]).unwrap()[0];
    let graph = builder.build();

    let public_plan = api.plan_public(&graph).unwrap();

    assert_eq!(
        public_plan.backend_policy,
        BackendPolicy::PreferGpuWithCpuFallback
    );
    assert_eq!(public_plan.requested_target.kind, DeviceKind::Gpu);
    #[cfg(not(any(feature = "accelerated-pilot", feature = "rocm-hip")))]
    {
        assert_eq!(public_plan.selected_target.kind, DeviceKind::Cpu);
        assert!(
            public_plan
                .fallback_reason
                .as_ref()
                .unwrap()
                .contains("gpu_scaffold has no executable kernel for add")
        );
    }
    let mut inputs = TensorStore::new();
    inputs.insert(
        lhs,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![4]),
            vec![1, 2, 3, 4],
        ),
    );
    inputs.insert(
        rhs,
        Tensor::dense_cpu(
            DomainId::new("integer"),
            Shape::from(vec![4]),
            vec![10, 20, 30, 40],
        ),
    );
    let (store, report) = api
        .execute_i64_public(&graph, &public_plan, inputs, &[out])
        .unwrap();
    assert_eq!(store.get(out).unwrap().data, vec![11, 22, 33, 44]);
    assert_eq!(report.checked_outputs, vec![out]);
}

#[test]
fn p190_public_api_strict_gpu_rejects_unsupported_nonstandard_domains() {
    let api = Tokitai::strict_gpu();
    let q5 = api.padic_domain(5, 3).unwrap();
    let err = api
        .plan_padic_sum_products(&q5, 0, 1, 2)
        .unwrap_err()
        .to_string();

    assert!(err.contains("strict GPU policy rejected plan"));
    assert!(err.contains("gpu_scaffold does not support fixed-precision p-adic execution"));
}

#[test]
fn p190_public_api_exposes_sheaf_gluing_with_policy_boundary() {
    let api = Tokitai::cpu_only();
    let cover = Cover::new("U", ["A", "B"]);
    let public_plan = api.plan_sheaf_glue_check(&cover).unwrap();
    assert_eq!(public_plan.selected_target.kind, DeviceKind::Cpu);

    let site = FiniteSite::new(
        vec![open("U"), open("A"), open("B"), open("A_cap_B")],
        vec![
            Inclusion::new("A", "U"),
            Inclusion::new("B", "U"),
            Inclusion::new("A_cap_B", "A"),
            Inclusion::new("A_cap_B", "B"),
            Inclusion::new("A_cap_B", "U"),
        ],
    )
    .with_intersection(open("A"), open("B"), open("A_cap_B"));
    let mut sections = SectionTable::new();
    sections.insert(open("A"), 7);
    sections.insert(open("B"), 7);
    sections.insert(open("A_cap_B"), 7);

    let glued = api
        .verify_sheaf_glue(&public_plan, &site, &cover, &sections, 42)
        .unwrap();
    assert_eq!(glued.open, open("U"));
    assert_eq!(glued.value, 42);
}

// P452: facade builder smoke coverage. Every Tokitai builder method
// (api.add(), api.softmax(), etc.) is exercised at least once through
// a planned graph. The test only asserts that the planner accepts the
// graph — it does not execute every branch. The point is to surface
// drift: if a builder method's signature changes or its op type is
// renamed, this test breaks at compile time.
#[test]
fn p452_facade_builder_methods_plan_successfully() {
    use tokitai_operator::op::arithmetic::{
        AbsOp, Exp2Op, Log2Op, MulByTwoOp, PowOp, ScalarAddOp, ScalarMulOp, SqrtOp, SquareOp, SubOp,
    };
    use tokitai_operator::op::index::{GatherOp, IndexAddOp, IndexSelectOp, NonzeroOp, ScatterOp};
    use tokitai_operator::op::nn::{GeluOp, LayerNormOp, ReluOp, SigmoidOp, TanhOp};
    use tokitai_operator::op::reductions::{
        AllOp, AnyOp, ArgMaxOp, ArgMinOp, MaxOp, MeanOp, MinOp, ProdOp, SumOp,
    };
    use tokitai_operator::op::shape::{ConcatOp, FlattenOp};
    use tokitai_operator::op::{MapOp, MatmulOp, ReduceOp};

    let api = Tokitai::cpu_only();
    let domain = DomainId::new("integer");
    let make_meta = |dims: Vec<usize>| {
        ObjectMeta::tensor(
            domain.clone(),
            Shape::from(dims),
            Representation::dense_cpu(),
        )
    };
    let v4 = make_meta(vec![4]);

    let mut builder = FacadeGraphBuilder::new();
    // Unary builders on a rank-1 input.
    let x = builder.input(v4.clone());
    let _ = builder.add_op(api.abs(), &[x]).unwrap();
    let _ = builder.add_op(api.neg(), &[x]).unwrap();
    let _ = builder.add_op(api.square(), &[x]).unwrap();
    let _ = builder.add_op(api.mul_by_two(), &[x]).unwrap();
    let _ = builder.add_op(AbsOp, &[x]).unwrap();
    let _ = builder.add_op(SquareOp, &[x]).unwrap();
    let _ = builder.add_op(MulByTwoOp, &[x]).unwrap();
    let _ = builder.add_op(SqrtOp, &[x]).unwrap();
    let _ = builder.add_op(Exp2Op, &[x]).unwrap();
    let _ = builder.add_op(Log2Op, &[x]).unwrap();
    let _ = builder.add_op(MapOp, &[x]).unwrap();
    let _ = builder.add_op(ReduceOp, &[x]).unwrap();
    let _ = builder.add_op(ReluOp, &[x]).unwrap();
    let _ = builder.add_op(SigmoidOp, &[x]).unwrap();
    let _ = builder.add_op(TanhOp, &[x]).unwrap();
    let _ = builder.add_op(GeluOp, &[x]).unwrap();
    let _ = builder.add_op(api.softmax(), &[x]).unwrap();
    // Reductions take (input, axis) — pass a second tensor for axis.
    let axis = builder.input(v4.clone());
    let _ = builder.add_op(SumOp, &[x, axis]).unwrap();
    let _ = builder.add_op(MeanOp, &[x, axis]).unwrap();
    let _ = builder.add_op(MaxOp, &[x, axis]).unwrap();
    let _ = builder.add_op(MinOp, &[x, axis]).unwrap();
    let _ = builder.add_op(ArgMaxOp, &[x, axis]).unwrap();
    let _ = builder.add_op(ArgMinOp, &[x, axis]).unwrap();
    let _ = builder.add_op(ProdOp, &[x, axis]).unwrap();
    let _ = builder.add_op(AnyOp, &[x, axis]).unwrap();
    let _ = builder.add_op(AllOp, &[x, axis]).unwrap();
    let _ = builder.add_op(NonzeroOp, &[x]).unwrap();

    // Binary builders on two same-shape rank-1 inputs.
    let a = builder.input(v4.clone());
    let b = builder.input(v4.clone());
    let _ = builder.add_op(api.add(), &[a, b]).unwrap();
    let _ = builder.add_op(api.mul(), &[a, b]).unwrap();
    let _ = builder.add_op(SubOp, &[a, b]).unwrap();
    let _ = builder.add_op(api.div(), &[a, b]).unwrap();
    let _ = builder.add_op(api.p_dot(), &[a, b]).unwrap();
    let _ = builder.add_op(ScalarAddOp, &[a, b]).unwrap();
    let _ = builder.add_op(ScalarMulOp, &[a, b]).unwrap();
    let _ = builder.add_op(PowOp, &[a, b]).unwrap();

    // Ternary builders: clamp (data, lo, hi), fma (a, b, c).
    let lo = builder.input(v4.clone());
    let hi = builder.input(v4.clone());
    let _ = builder.add_op(api.clamp(), &[a, lo, hi]).unwrap();
    let _ = builder.add_op(api.fma(), &[a, b, b]).unwrap();

    // Matmul on two rank-2 inputs.
    let m1 = builder.input(make_meta(vec![2, 3]));
    let m2 = builder.input(make_meta(vec![3, 2]));
    let _ = builder.add_op(MatmulOp, &[m1, m2]).unwrap();

    // Layer norm needs (input, gamma, beta) — pass the same tensor for gamma/beta.
    let ln = builder.input(make_meta(vec![2, 3]));
    let _ = builder.add_op(LayerNormOp, &[ln, ln, ln]).unwrap();

    // Concat + Flatten on rank-1 / rank-2 inputs respectively.
    let _ = builder.add_op(ConcatOp, &[a, b]).unwrap();
    let _ = builder.add_op(FlattenOp, &[ln]).unwrap();

    // Index builders: gather (input, indices, axis), scatter (input,
    // indices, updates), index_select (input, indices, axis), index_add
    // (input, indices, updates, axis).
    let idx = builder.input(v4.clone());
    let _ = builder.add_op(GatherOp, &[a, idx, idx]).unwrap();
    let _ = builder.add_op(ScatterOp, &[a, idx, b]).unwrap();
    let _ = builder.add_op(IndexSelectOp, &[a, idx, idx]).unwrap();
    let _ = builder.add_op(IndexAddOp, &[a, idx, b, idx]).unwrap();

    let graph = builder.build();
    // Smoke assertion: the planner accepts the graph (every node
    // produced a valid infer output). Per-op execution correctness
    // is covered by the per-op test files.
    let _plan = api
        .plan(&graph)
        .expect("planner should accept facade-built graph");
}