// Copyright (c) 2022-2022, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

// The original work for this formula was implmented in aomenc, and this is
// an adaptation of that work:
// https://aomedia.googlesource.com/aom/+/refs/heads/main/examples/photon_noise_table.c

// This implementation creates a film grain table, for use in stills and videos,
// representing the noise that one would get by shooting with a digital camera
// at a given light level. Much of the noise in digital images is photon shot
// noise, which is due to the characteristics of photon arrival and grows in
// standard deviation as the square root of the expected number of photons
// captured.
// https://www.photonstophotos.net/Emil%20Martinec/noise.html#shotnoise
//
// The proxy used by this implementation for the amount of light captured
// is the ISO value such that the focal plane exposure at the time of capture
// would have been mapped by a 35mm camera to the output lightness observed
// in the image. That is, if one were to shoot on a 35mm camera (36×24mm sensor)
// at the nominal exposure for that ISO setting, the resulting image should
// contain noise of the same order of magnitude as generated by this
// implementation.
//
// The (mostly) square-root relationship between light intensity and noise
// amplitude holds in linear light, but AV1 streams are most often encoded
// non-linearly, and the film grain is applied to those non-linear values.
// Therefore, this implementation must account for the non-linearity, and this
// is controlled by the transfer function parameter, which specifies the tone
// response curve that will be used when encoding the actual image. The default
// for this implementation is BT.1886, which is approximately similar to an
// encoding gamma of 1/2.8 (i.e. a decoding gamma of 2.8) though not quite
// identical.
//
// As alluded to above, the implementation assumes that the image is taken from
// the entirety of a 36×24mm (“35mm format”) sensor. If that assumption does not
// hold, then a “35mm-equivalent ISO value” that can be passed to the
// implementation can be obtained by multiplying the true ISO value by the ratio
// of 36×24mm to the area that was actually used. For formats that approximately
// share the same aspect ratio, this is often expressed as the square of the
// “equivalence ratio” which is the ratio of their diagonals. For example, APS-C
// (often ~24×16mm) is said to have an equivalence ratio of 1.5 relative to the
// 35mm format, and therefore ISO 1000 on APS-C and ISO 1000×1.5² = 2250 on 35mm
// produce an image of the same lightness from the same amount of light spread
// onto their respective surface areas (resulting in different focal plane
// exposures), and those images will thus have similar amounts of noise if the
// cameras are of similar technology. https://doi.org/10.1117/1.OE.57.11.110801
//
// The implementation needs to know the resolution of the images to which its
// grain tables will be applied so that it can know how the light on the sensor
// was shared between its pixels. As a general rule, while a higher pixel count
// will lead to more noise per pixel, when the final image is viewed at the same
// physical size, that noise will tend to “average out” to the same amount over
// a given area, since there will be more pixels in it which, in aggregate, will
// have received essentially as much light. Put differently, the amount of noise
// depends on the scale at which it is measured, and the decision for this
// implementation was to make that scale relative to the image instead of its
// constituent samples. For more on this, see:
//
// https://www.photonstophotos.net/Emil%20Martinec/noise-p3.html#pixelsize
// https://www.dpreview.com/articles/5365920428/the-effect-of-pixel-and-sensor-sizes-on-noise/2
// https://www.dpreview.com/videos/7940373140/dpreview-tv-why-lower-resolution-sensors-are-not-better-in-low-light

use std::{
    fs::File,
    io::{BufWriter, Write},
    path::Path,
};

use arrayvec::ArrayVec;

use crate::{GrainTableSegment, ScalingPoints, DEFAULT_GRAIN_SEED, NUM_Y_POINTS};

const PQ_M1: f32 = 2610. / 16384.;
const PQ_M2: f32 = 128. * 2523. / 4096.;
const PQ_C1: f32 = 3424. / 4096.;
const PQ_C2: f32 = 32. * 2413. / 4096.;
const PQ_C3: f32 = 32. * 2392. / 4096.;

const BT1886_WHITEPOINT: f32 = 203.;
const BT1886_BLACKPOINT: f32 = 0.1;
const BT1886_GAMMA: f32 = 2.4;

// BT.1886 formula from https://en.wikipedia.org/wiki/ITU-R_BT.1886.
//
// TODO: the inverses, alpha, and beta should all be constants
// once floats in const fns are stabilized and `powf` is const.
// Until then, `inline(always)` gets us close enough.

#[inline(always)]
fn bt1886_inv_whitepoint() -> f32 {
    BT1886_WHITEPOINT.powf(1.0 / BT1886_GAMMA)
}

#[inline(always)]
fn bt1886_inv_blackpoint() -> f32 {
    BT1886_BLACKPOINT.powf(1.0 / BT1886_GAMMA)
}

/// The variable for user gain:
/// `α = (Lw^(1/λ) - Lb^(1/λ)) ^ λ`
#[inline(always)]
fn bt1886_alpha() -> f32 {
    (bt1886_inv_whitepoint() - bt1886_inv_blackpoint()).powf(BT1886_GAMMA)
}

/// The variable for user black level lift:
/// `β = Lb^(1/λ) / (Lw^(1/λ) - Lb^(1/λ))`
#[inline(always)]
fn bt1886_beta() -> f32 {
    bt1886_inv_blackpoint() / (bt1886_inv_whitepoint() - bt1886_inv_blackpoint())
}

/// Settings and video data defining how to generate the film grain params.
#[derive(Debug, Clone, Copy)]
pub struct NoiseGenArgs {
    pub iso_setting: u32,
    pub width: u32,
    pub height: u32,
    pub transfer_function: TransferFunction,
    pub chroma_grain: bool,
    pub random_seed: Option<u16>,
}

/// Generates a set of photon noise parameters for a segment of video
/// given a set of `args`.
#[must_use]
pub fn generate_photon_noise_params(
    start_time: u64,
    end_time: u64,
    args: NoiseGenArgs,
) -> GrainTableSegment {
    GrainTableSegment {
        start_time,
        end_time,
        scaling_points_y: generate_luma_noise_points(args),
        scaling_points_cb: ArrayVec::new(),
        scaling_points_cr: ArrayVec::new(),
        scaling_shift: 8,
        ar_coeff_lag: 0,
        ar_coeffs_y: ArrayVec::new(),
        ar_coeffs_cb: ArrayVec::try_from([0].as_slice())
            .expect("Cannot fail creation from const array"),
        ar_coeffs_cr: ArrayVec::try_from([0].as_slice())
            .expect("Cannot fail creation from const array"),
        ar_coeff_shift: 6,
        cb_mult: 0,
        cb_luma_mult: 0,
        cb_offset: 0,
        cr_mult: 0,
        cr_luma_mult: 0,
        cr_offset: 0,
        overlap_flag: true,
        chroma_scaling_from_luma: args.chroma_grain,
        grain_scale_shift: 0,
        random_seed: args.random_seed.unwrap_or(DEFAULT_GRAIN_SEED),
    }
}

/// Generates a set of film grain parameters for a segment of video
/// given a set of `args`.
///
/// # Panics
/// - This is not yet implemented, so it will always panic
#[must_use]
#[cfg(feature = "unstable")]
pub fn generate_film_grain_params(
    start_time: u64,
    end_time: u64,
    args: NoiseGenArgs,
) -> GrainTableSegment {
    todo!("SCIENCE");
    // GrainTableSegment {
    //     start_time,
    //     end_time,
    //     scaling_points_y: generate_luma_noise_points(args),
    //     scaling_points_cb: ArrayVec::new(),
    //     scaling_points_cr: ArrayVec::new(),
    //     scaling_shift: 8,
    //     ar_coeff_lag: 0,
    //     ar_coeffs_y: ArrayVec::new(),
    //     ar_coeffs_cb: ArrayVec::try_from([0].as_slice())
    //         .expect("Cannot fail creation from const array"),
    //     ar_coeffs_cr: ArrayVec::try_from([0].as_slice())
    //         .expect("Cannot fail creation from const array"),
    //     ar_coeff_shift: 6,
    //     cb_mult: 0,
    //     cb_luma_mult: 0,
    //     cb_offset: 0,
    //     cr_mult: 0,
    //     cr_luma_mult: 0,
    //     cr_offset: 0,
    //     overlap_flag: true,
    //     chroma_scaling_from_luma: args.chroma_grain,
    //     grain_scale_shift: 0,
    //     random_seed: args.random_seed.unwrap_or(DEFAULT_GRAIN_SEED),
    // }
}

/// Write a set of generated film grain params to a table file,
/// using the standard film grain table format supported by
/// aomenc, rav1e, and svt-av1.
///
/// # Errors
///
/// - If the output file cannot be written to
pub fn write_grain_table<P: AsRef<Path>>(
    filename: P,
    params: &[GrainTableSegment],
) -> anyhow::Result<()> {
    let mut file = BufWriter::new(File::create(filename)?);
    writeln!(&mut file, "filmgrn1")?;
    for segment in params {
        write_film_grain_segment(segment, &mut file)?;
    }
    file.flush()?;

    Ok(())
}

fn write_film_grain_segment(
    params: &GrainTableSegment,
    output: &mut BufWriter<File>,
) -> anyhow::Result<()> {
    writeln!(
        output,
        "E {} {} 1 {} 1",
        params.start_time, params.end_time, params.random_seed,
    )?;
    writeln!(
        output,
        "\tp {} {} {} {} {} {} {} {} {} {} {} {}",
        params.ar_coeff_lag,
        params.ar_coeff_shift,
        params.grain_scale_shift,
        params.scaling_shift,
        u8::from(params.chroma_scaling_from_luma),
        u8::from(params.overlap_flag),
        params.cb_mult,
        params.cb_luma_mult,
        params.cb_offset,
        params.cr_mult,
        params.cr_luma_mult,
        params.cr_offset
    )?;

    write!(output, "\tsY {} ", params.scaling_points_y.len())?;
    for point in &params.scaling_points_y {
        write!(output, " {} {}", point[0], point[1])?;
    }
    writeln!(output)?;

    write!(output, "\tsCb {}", params.scaling_points_cb.len())?;
    for point in &params.scaling_points_cb {
        write!(output, " {} {}", point[0], point[1])?;
    }
    writeln!(output)?;

    write!(output, "\tsCr {}", params.scaling_points_cr.len())?;
    for point in &params.scaling_points_cr {
        write!(output, " {} {}", point[0], point[1])?;
    }
    writeln!(output)?;

    write!(output, "\tcY")?;
    for coeff in &params.ar_coeffs_y {
        write!(output, " {}", *coeff)?;
    }
    writeln!(output)?;

    write!(output, "\tcCb")?;
    for coeff in &params.ar_coeffs_cb {
        write!(output, " {}", *coeff)?;
    }
    writeln!(output)?;

    write!(output, "\tcCr")?;
    for coeff in &params.ar_coeffs_cr {
        write!(output, " {}", *coeff)?;
    }
    writeln!(output)?;

    Ok(())
}

#[allow(clippy::upper_case_acronyms)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransferFunction {
    /// For SDR content
    BT1886,
    /// For HDR content
    SMPTE2084,
}

impl TransferFunction {
    #[must_use]
    pub fn to_linear(self, x: f32) -> f32 {
        match self {
            TransferFunction::BT1886 => {
                // The screen luminance in cd/m^2:
                // L = α * (x + β)^λ
                let luma = bt1886_alpha() * (x + bt1886_beta()).powf(BT1886_GAMMA);

                // Normalize to between 0.0 and 1.0
                luma / BT1886_WHITEPOINT
            }
            TransferFunction::SMPTE2084 => {
                let pq_pow_inv_m2 = x.powf(1. / PQ_M2);
                (0_f32.max(pq_pow_inv_m2 - PQ_C1) / (PQ_C2 - PQ_C3 * pq_pow_inv_m2))
                    .powf(1. / PQ_M1)
            }
        }
    }

    #[allow(clippy::wrong_self_convention)]
    #[must_use]
    pub fn from_linear(self, x: f32) -> f32 {
        match self {
            TransferFunction::BT1886 => {
                // Scale to a raw cd/m^2 value
                let luma = x * BT1886_WHITEPOINT;

                // The inverse of the `to_linear` formula:
                // `(L / α)^(1 / λ) - β = x`
                (luma / bt1886_alpha()).powf(1.0 / BT1886_GAMMA) - bt1886_beta()
            }
            TransferFunction::SMPTE2084 => {
                if x < f32::EPSILON {
                    return 0.0;
                }
                let linear_pow_m1 = x.powf(PQ_M1);
                (PQ_C2.mul_add(linear_pow_m1, PQ_C1) / PQ_C3.mul_add(linear_pow_m1, 1.)).powf(PQ_M2)
            }
        }
    }

    #[inline(always)]
    #[must_use]
    pub fn mid_tone(self) -> f32 {
        self.to_linear(0.5)
    }
}

fn generate_luma_noise_points(args: NoiseGenArgs) -> ScalingPoints {
    // Assumes a daylight-like spectrum.
    // https://www.strollswithmydog.com/effective-quantum-efficiency-of-sensor/#:~:text=11%2C260%20photons/um%5E2/lx-s
    const PHOTONS_PER_SQ_MICRON_PER_LUX_SECOND: f32 = 11260.;

    // Order of magnitude for cameras in the 2010-2020 decade, taking the CFA into
    // account.
    const EFFECTIVE_QUANTUM_EFFICIENCY: f32 = 0.2;

    // Also reasonable values for current cameras. The read noise is typically
    // higher than this at low ISO settings but it matters less there.
    const PHOTO_RESPONSE_NON_UNIFORMITY: f32 = 0.005;
    const INPUT_REFERRED_READ_NOISE: f32 = 1.5;

    // Assumes a 35mm sensor (36mm × 24mm).
    const SENSOR_AREA: f32 = 36_000. * 24_000.;

    // Focal plane exposure for a mid-tone (typically a 18% reflectance card), in
    // lx·s.
    let mid_tone_exposure = 10. / args.iso_setting as f32;

    let pixel_area_microns = SENSOR_AREA / (args.width * args.height) as f32;

    let mid_tone_electrons_per_pixel = EFFECTIVE_QUANTUM_EFFICIENCY
        * PHOTONS_PER_SQ_MICRON_PER_LUX_SECOND
        * mid_tone_exposure
        * pixel_area_microns;
    let max_electrons_per_pixel = mid_tone_electrons_per_pixel / args.transfer_function.mid_tone();

    let mut scaling_points = ScalingPoints::default();
    for i in 0..NUM_Y_POINTS {
        let x = i as f32 / (NUM_Y_POINTS as f32 - 1.);
        let linear = args.transfer_function.to_linear(x);
        let electrons_per_pixel = max_electrons_per_pixel * linear;

        // Quadrature sum of the relevant sources of noise, in electrons rms. Photon
        // shot noise is sqrt(electrons) so we can skip the square root and the
        // squaring.
        // https://en.wikipedia.org/wiki/Addition_in_quadrature
        // https://doi.org/10.1117/3.725073
        let noise_in_electrons = (PHOTO_RESPONSE_NON_UNIFORMITY
            * PHOTO_RESPONSE_NON_UNIFORMITY
            * electrons_per_pixel)
            .mul_add(
                electrons_per_pixel,
                INPUT_REFERRED_READ_NOISE.mul_add(INPUT_REFERRED_READ_NOISE, electrons_per_pixel),
            )
            .sqrt();
        let linear_noise = noise_in_electrons / max_electrons_per_pixel;
        let linear_range_start = 0_f32.max(linear - 2. * linear_noise);
        let linear_range_end = 1_f32.min(2_f32.mul_add(linear_noise, linear));
        let tf_slope = (args.transfer_function.from_linear(linear_range_end)
            - args.transfer_function.from_linear(linear_range_start))
            / (linear_range_end - linear_range_start);
        let encoded_noise = linear_noise * tf_slope;

        let x = (255. * x).round() as u8;
        let encoded_noise = 255_f32.min((255. * 7.88 * encoded_noise).round()) as u8;

        scaling_points.push([x, encoded_noise]);
    }

    scaling_points
}

#[cfg(test)]
mod tests {
    use quickcheck::TestResult;
    use quickcheck_macros::quickcheck;

    use super::*;

    #[quickcheck]
    fn bt1886_to_linear_within_range(x: f32) -> TestResult {
        if !(0.0..=1.0).contains(&x) || x.is_nan() {
            return TestResult::discard();
        }

        let tx = TransferFunction::BT1886;
        let res = tx.to_linear(x);
        TestResult::from_bool((0.0..=1.0).contains(&res))
    }

    #[quickcheck]
    fn bt1886_to_linear_reverts_correctly(x: f32) -> TestResult {
        if !(0.0..=1.0).contains(&x) || x.is_nan() {
            return TestResult::discard();
        }

        let tx = TransferFunction::BT1886;
        let res = tx.to_linear(x);
        let res = tx.from_linear(res);
        TestResult::from_bool((x - res).abs() < f32::EPSILON)
    }

    #[quickcheck]
    fn smpte2084_to_linear_within_range(x: f32) -> TestResult {
        if !(0.0..=1.0).contains(&x) || x.is_nan() {
            return TestResult::discard();
        }

        let tx = TransferFunction::SMPTE2084;
        let res = tx.to_linear(x);
        TestResult::from_bool((0.0..=1.0).contains(&res))
    }

    #[quickcheck]
    fn smpte2084_to_linear_reverts_correctly(x: f32) -> TestResult {
        if !(0.0..=1.0).contains(&x) || x.is_nan() {
            return TestResult::discard();
        }

        let tx = TransferFunction::SMPTE2084;
        let res = tx.to_linear(x);
        let res = tx.from_linear(res);
        TestResult::from_bool((x - res).abs() < f32::EPSILON)
    }
}