#include "tests.h"

#include <math.h>

void run_noise_tests(void)
{
    psi_test_section("Noise channels");

    struct PsiDensityMatrix pure = psi_new_density_matrix(1);
    psi_test_check(psi_is_pure_density_matrix(pure, 1e-10), "fresh density matrix is pure");
    psi_test_check(fabs(psi_trace_density_matrix(pure).real - 1.0) < 1e-10, "trace = 1");
    psi_free_density_matrix(&pure);

    struct PsiDensityMatrix bf = psi_new_density_matrix(1);
    struct PsiNoiseChannel flip = psi_bit_flip_channel(0.25);
    psi_apply_noise_channel(&bf, flip, 0);
    double probs[2];
    psi_density_matrix_probabilities(bf, probs);
    psi_test_check(fabs(probs[0] - 0.75) < 1e-10 && fabs(probs[1] - 0.25) < 1e-10,
                   "bit_flip(0.25) on |0>: p = [0.75, 0.25]");
    psi_test_check(fabs(psi_purity_density_matrix(bf) - 0.625) < 1e-10,
                   "bit_flip(0.25) purity = 0.625");
    psi_free_noise_channel(&flip);
    psi_free_density_matrix(&bf);

    struct PsiComplex one[2] = { psi_new_complex(0.0, 0.0), psi_new_complex(1.0, 0.0) };
    struct PsiDensityMatrix damp = psi_new_density_matrix_from_state(one, 2);
    struct PsiNoiseChannel ad = psi_amplitude_damping_channel(1.0);
    psi_apply_noise_channel(&damp, ad, 0);
    double dprobs[2];
    psi_density_matrix_probabilities(damp, dprobs);
    psi_test_check(fabs(dprobs[0] - 1.0) < 1e-10, "amplitude_damping(1.0) fully decays |1> -> |0>");
    psi_free_noise_channel(&ad);
    psi_free_density_matrix(&damp);

    struct PsiComplex zero[2] = { psi_new_complex(1.0, 0.0), psi_new_complex(0.0, 0.0) };
    struct PsiDensityMatrix fid = psi_new_density_matrix_from_state(zero, 2);
    psi_test_check(fabs(psi_fidelity_density_matrix(fid, zero) - 1.0) < 1e-10,
                   "fidelity(|0><0|, |0>) = 1");
    psi_free_density_matrix(&fid);

    struct PsiDensityMatrix depo = psi_new_density_matrix(1);
    struct PsiNoiseChannel dc = psi_depolarising_channel(0.3);
    psi_apply_noise_channel(&depo, dc, 0);
    psi_test_check(fabs(psi_trace_density_matrix(depo).real - 1.0) < 1e-10,
                   "depolarising preserves trace");
    psi_free_noise_channel(&dc);
    psi_free_density_matrix(&depo);
}