// 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.

fn bt1886_inv_whitepoint() -> f32 {

    BT1886_WHITEPOINT.powf(1.0 / BT1886_GAMMA)

}

fn bt1886_inv_blackpoint() -> f32 {

    BT1886_BLACKPOINT.powf(1.0 / BT1886_GAMMA)

}

/// The variable for user gain:

/// `α = (Lw^(1/λ) - Lb^(1/λ)) ^ λ`

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/λ))`

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]

#[inline]

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]

#[inline]

#[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

#[inline]

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]

    #[inline]

    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_C3.mul_add(-pq_pow_inv_m2, PQ_C2))

                    .powf(1. / PQ_M1)

            }

        }

    }

    #[allow(clippy::wrong_self_convention)]

    #[must_use]

    #[inline]

    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]

    #[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(2.0f32.mul_add(-linear_noise, linear));

        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)

    }

}