use std::{fs, mem, path::Path};

use color_eyre::eyre::{bail, Result};
use image::{ImageBuffer, Luma, Pixel, Rgba, RgbaImage};
use nalgebra::{vector, Vector2};
use num_traits::identities::Zero;

use super::{
    gpx::Coordinates,
    layer::{self, TileLayer},
};

pub type HeatCounter = TileLayer<Luma<u32>>;

pub type HeatMap = TileLayer<Rgba<u8>>;

/// Returns (a - b)**2, but ensures that no underflow happens (if b > a).
fn diff_squared(a: u64, b: u64) -> u64 {
    if a > b {
        (a - b).pow(2)
    } else {
        (b - a).pow(2)
    }
}

fn render_circle<P: Pixel>(layer: &mut TileLayer<P>, center: (u64, u64), radius: u64, pixel: P) {
    let x_lower = center.0.saturating_sub(radius);
    let x_upper = (layer.width() - 1).min(center.0 + radius);
    let y_lower = center.1.saturating_sub(radius);
    let y_upper = (layer.height() - 1).min(center.1 + radius);

    for x in x_lower..=x_upper {
        for y in y_lower..=y_upper {
            if diff_squared(center.0, x) + diff_squared(center.1, y) <= radius * radius {
                *layer.get_pixel_mut(x, y) = pixel;
            }
        }
    }
}

fn direction_vector(a: (u64, u64), b: (u64, u64)) -> Vector2<f64> {
    let dx = if b.0 > a.0 {
        (b.0 - a.0) as f64
    } else {
        -((a.0 - b.0) as f64)
    };
    let dy = if b.1 > a.1 {
        (b.1 - a.1) as f64
    } else {
        -((a.1 - b.1) as f64)
    };
    vector![dx, dy]
}

fn render_line<P: Pixel>(
    layer: &mut TileLayer<P>,
    start: (u64, u64),
    end: (u64, u64),
    thickness: u64,
    pixel: P,
) {
    use imageproc::point::Point;

    if start == end {
        return;
    }

    fn unsigned_add(a: Vector2<u64>, b: Vector2<i32>) -> Vector2<u64> {
        let x = if b[0] < 0 {
            a[0] - b[0].abs() as u64
        } else {
            a[0] + b[0] as u64
        };
        let y = if b[1] < 0 {
            a[1] - b[1].abs() as u64
        } else {
            a[1] + b[1] as u64
        };
        vector![x, y]
    }

    let r = direction_vector(start, end);
    let normal = vector![r[1], -r[0]].normalize();

    let start = vector![start.0, start.1];
    let end = vector![end.0, end.1];

    let displacement = normal * thickness as f64;
    let displacement = displacement.map(|x| x as i32);
    let polygon = [
        unsigned_add(start, displacement),
        unsigned_add(end, displacement),
        unsigned_add(end, -displacement),
        unsigned_add(start, -displacement),
    ];
    let min_x = polygon.iter().map(|p| p[0]).min().unwrap();
    let min_y = polygon.iter().map(|p| p[1]).min().unwrap();
    let max_x = polygon.iter().map(|p| p[0]).max().unwrap();
    let max_y = polygon.iter().map(|p| p[1]).max().unwrap();

    let mut overlay = ImageBuffer::<P, Vec<P::Subpixel>>::new(
        (max_x - min_x).try_into().unwrap(),
        (max_y - min_y).try_into().unwrap(),
    );
    let adjusted_poly = polygon
        .into_iter()
        .map(|p| Point::new((p[0] - min_x) as i32, (p[1] - min_y) as i32))
        .collect::<Vec<_>>();
    imageproc::drawing::draw_polygon_mut(&mut overlay, &adjusted_poly, pixel);

    for (x, y, pixel) in overlay.enumerate_pixels() {
        if pixel.channels()[0] > Zero::zero() {
            *layer.get_pixel_mut(u64::from(x) + min_x, u64::from(y) + min_y) = *pixel;
        }
    }
}

fn merge_heat_counter(base: &mut HeatCounter, overlay: &HeatCounter) {
    for (x, y, source) in overlay.enumerate_pixels() {
        let target = base.get_pixel_mut(x, y);
        target[0] += source[0];
    }
}

fn colorize_tile(tile: &ImageBuffer<Luma<u32>, Vec<u32>>, max: u32) -> RgbaImage {
    let gradient = colorgrad::turbo();
    let mut result = ImageBuffer::from_pixel(tile.width(), tile.height(), [0, 0, 0, 0].into());
    for (x, y, pixel) in tile.enumerate_pixels() {
        if pixel[0] > 0 {
            let alpha = pixel[0] as f64 / max as f64;
            let color = gradient.at(alpha);
            let target = result.get_pixel_mut(x, y);
            *target = color.to_rgba8().into();
        }
    }
    result
}

pub fn colorize_heatcounter(layer: &HeatCounter) -> HeatMap {
    let max = layer.pixels().map(|l| l.0[0]).max().unwrap_or_default();
    let mut result = TileLayer::from_pixel(layer.width(), layer.height(), [0, 0, 0, 0].into());
    if max == 0 {
        return result;
    }
    for (tile_x, tile_y, tile) in layer.enumerate_tiles() {
        let colorized = colorize_tile(&tile, max);
        *result.tile_mut(tile_x, tile_y) = colorized;
    }
    result
}

/// Lazily colorizes a [`HeatCounter`] by colorizing it tile-by-tile and saving a tile before
/// rendering the next one.
///
/// This has a way lower memory usage than [`colorize_heatcounter`].
pub fn lazy_colorization<P: AsRef<Path>>(layer: &HeatCounter, base_dir: P) -> Result<()> {
    let base_dir = base_dir.as_ref();
    let max = layer.pixels().map(|l| l.0[0]).max().unwrap_or_default();
    if max == 0 {
        return Ok(());
    }
    for (tile_x, tile_y, tile) in layer.enumerate_tiles() {
        let colorized = colorize_tile(&tile, max);
        let folder = base_dir.join(&tile_x.to_string());
        let metadata = folder.metadata();
        match metadata {
            Err(_) => fs::create_dir(&folder)?,
            Ok(m) if !m.is_dir() => bail!("Output path is not a directory"),
            _ => {}
        }
        let file = folder.join(&format!("{tile_y}.png"));
        layer::compress_png(&colorized, file)?;
    }
    Ok(())
}

/// Renders the heat counter for the given zoom level and track points.
pub fn render_heatcounter(zoom: u32, tracks: &[Vec<Coordinates>]) -> HeatCounter {
    let size = 256 * 2u64.pow(zoom);

    let mut heatcounter = TileLayer::from_pixel(size, size, [0].into());

    for track in tracks {
        let mut layer = TileLayer::from_pixel(size, size, [0].into());

        let points = track
            .iter()
            .map(|coords| coords.web_mercator(zoom))
            .collect::<Vec<_>>();

        for point in points.iter() {
            render_circle(&mut layer, *point, 10, [1].into());
        }

        for (a, b) in points.iter().zip(points.iter().skip(1)) {
            render_line(&mut layer, *a, *b, 10, [1].into());
        }

        merge_heat_counter(&mut heatcounter, &layer);
    }
    heatcounter
}