// -*- coding: utf-8 -*-
//
// Copyright 2021-2023 Michael Büsch <m@bues.ch>
//
// Licensed under the Apache License version 2.0
// or the MIT license, at your option.
// SPDX-License-Identifier: Apache-2.0 OR MIT
//

use anyhow as ah;
use colorsys::{ColorAlpha, ColorTransform, Rgb};
use rand::{distributions::Standard, prelude::*};
use rand_chacha::ChaCha8Rng;
use std::path::{Path, PathBuf};

const ZSHIFT: usize = 8;

fn linear_interp(a: u32, b: u32, left: usize, center: usize, right: usize) -> u32 {
    ((a as usize * (right - center)) / (right - left)
        + (b as usize * (center - left)) / (right - left)) as u32
}

fn linear_interp_argb(a: u32, b: u32, left: usize, center: usize, right: usize) -> u32 {
    let ca = linear_interp(
        (a & 0xFF000000) >> 24,
        (b & 0xFF000000) >> 24,
        left,
        center,
        right,
    );
    let cr = linear_interp(
        (a & 0x00FF0000) >> 16,
        (b & 0x00FF0000) >> 16,
        left,
        center,
        right,
    );
    let cg = linear_interp(
        (a & 0x0000FF00) >> 8,
        (b & 0x0000FF00) >> 8,
        left,
        center,
        right,
    );
    let cb = linear_interp(a & 0x000000FF, b & 0x000000FF, left, center, right);

    ((ca & 0xFF) << 24) | ((cr & 0xFF) << 16) | ((cg & 0xFF) << 8) | (cb & 0xFF)
}

#[derive(Clone, PartialEq, Debug)]
pub struct MapValue {
    argb: u32,
    height: u16,
}

impl MapValue {
    fn new() -> MapValue {
        MapValue {
            argb: 0xFF000000,
            height: 0,
        }
    }

    #[inline]
    pub fn argb(&self) -> u32 {
        self.argb
    }

    #[inline]
    pub fn height(&self) -> u16 {
        self.height
    }
}

#[derive(Clone, PartialEq, Debug)]
#[repr(C, align(64))]
pub struct Map {
    data: Vec<MapValue>,
    x_size: usize,
    y_size: usize,
    color_map_file: Option<PathBuf>,
    height_map_file: Option<PathBuf>,
    height_step_thres: i32,
}

impl Map {
    pub const MAX_HEIGHT: u16 = u16::MAX;

    pub fn open_by_name(base_dir: &Path, name: &str) -> ah::Result<Map> {
        Self::open_by_name_and_suffix(base_dir, name, ".png", "h.png")
    }

    pub fn open_by_name_and_suffix(
        base_dir: &Path,
        name: &str,
        color_suffix: &str,
        height_suffix: &str,
    ) -> ah::Result<Map> {
        if base_dir.is_dir() {
            let mut color_map = base_dir.to_path_buf();
            color_map.push(name.to_string() + color_suffix);

            let mut height_map = base_dir.to_path_buf();
            height_map.push(name.to_string() + height_suffix);

            Self::open_files(color_map.as_path(), height_map.as_path())
        } else {
            Err(ah::format_err!(
                "Map: base_dir {:?} is not a directory.",
                base_dir
            ))
        }
    }

    pub fn open_files(color_map_file: &Path, height_map_file: &Path) -> ah::Result<Map> {
        let color_map = image::open(color_map_file)?.to_rgb8();
        let height_map = image::open(height_map_file)?.to_luma8();

        let x_size = color_map.width();
        let y_size = color_map.height();
        let secondary_x_size = height_map.width();
        let secondary_y_size = height_map.height();
        let height_div_x = x_size / secondary_x_size;
        let height_div_y = y_size / secondary_y_size;

        let mut data = vec![MapValue::new(); (x_size * y_size) as usize];

        for y in 0..y_size {
            for x in 0..x_size {
                let rgb = color_map.get_pixel(x, y);
                let argb =
                    0xFF << 24 | ((rgb[2] as u32) << 16) | ((rgb[1] as u32) << 8) | (rgb[0] as u32);
                let height = height_map.get_pixel(x / height_div_x, y / height_div_y)[0] as u16;

                let i = (y * x_size) + x;
                data[i as usize] = MapValue {
                    argb,
                    height: height << ZSHIFT,
                };
            }
        }

        Ok(Map::new(
            data,
            x_size as usize,
            y_size as usize,
            Some(color_map_file.to_path_buf()),
            Some(height_map_file.to_path_buf()),
        ))
    }

    fn new(
        data: Vec<MapValue>,
        x_size: usize,
        y_size: usize,
        color_map_file: Option<PathBuf>,
        height_map_file: Option<PathBuf>,
    ) -> Map {
        assert_eq!(data.len(), x_size * y_size);
        Map {
            data,
            x_size,
            y_size,
            color_map_file,
            height_map_file,
            height_step_thres: (1 << ZSHIFT) * 4,
        }
    }

    pub fn write_files(&self) -> ah::Result<()> {
        match (&self.color_map_file, &self.height_map_file) {
            (Some(color_map_file), Some(height_map_file)) => {
                let mut color_img = image::RgbImage::new(self.x_size as u32, self.y_size as u32);
                let mut height_img = image::RgbImage::new(self.x_size as u32, self.y_size as u32);

                for y in 0..self.y_size {
                    for x in 0..self.x_size {
                        let map_value = self.get(x, y);
                        let rgb = map_value.argb();
                        let height = map_value.height() >> ZSHIFT;

                        color_img.put_pixel(
                            x as u32,
                            y as u32,
                            [
                                (rgb & 0xFF) as u8,
                                ((rgb >> 8) & 0xFF) as u8,
                                ((rgb >> 16) & 0xFF) as u8,
                            ]
                            .into(),
                        );
                        height_img.put_pixel(
                            x as u32,
                            y as u32,
                            [
                                (height & 0xFF) as u8,
                                (height & 0xFF) as u8,
                                (height & 0xFF) as u8,
                            ]
                            .into(),
                        );
                    }
                }

                color_img.save_with_format(color_map_file, image::ImageFormat::Png)?;
                height_img.save_with_format(height_map_file, image::ImageFormat::Png)?;
                Ok(())
            }
            _ => Err(ah::format_err!("Map write_files: No file names set.")),
        }
    }

    #[inline]
    pub fn get(&self, x: usize, y: usize) -> &MapValue {
        let x = x % self.x_size;
        let y = y % self.y_size;
        let i = (y * self.x_size) + x;
        &self.data[i]
    }

    pub fn get_x_size(&self) -> usize {
        self.x_size
    }

    pub fn get_y_size(&self) -> usize {
        self.y_size
    }

    pub fn get_height_shift(&self) -> usize {
        ZSHIFT
    }

    #[inline]
    fn is_height_step(
        &self,
        q11: &MapValue,
        q12: &MapValue,
        q21: &MapValue,
        q22: &MapValue,
    ) -> bool {
        let height_diff0 = q11.height() as i32 - q22.height() as i32;
        let height_diff1 = q12.height() as i32 - q21.height() as i32;

        height_diff0.abs() > self.height_step_thres || height_diff1.abs() > self.height_step_thres
    }

    fn bilinear_interp_height(
        &mut self,
        x: usize,
        y: usize,
        x0: usize,
        y0: usize,
        x1: usize,
        y1: usize,
    ) {
        let q11 = self.get(x0, y0);
        let q12 = self.get(x0, y1);
        let q21 = self.get(x1, y0);
        let q22 = self.get(x1, y1);

        let new_height = {
            if self.is_height_step(q11, q12, q21, q22) {
                q11.height
            } else {
                let height_1 = linear_interp(q11.height() as u32, q21.height() as u32, x0, x, x1);
                let height_2 = linear_interp(q12.height() as u32, q22.height() as u32, x0, x, x1);
                linear_interp(height_1, height_2, y0, y, y1) as u16
            }
        };

        self.data[(y * self.x_size) + x].height = new_height;
    }

    fn bilinear_interp_color(
        &mut self,
        x: usize,
        y: usize,
        x0: usize,
        y0: usize,
        x1: usize,
        y1: usize,
    ) {
        let q11 = self.get(x0, y0);
        let q12 = self.get(x0, y1);
        let q21 = self.get(x1, y0);
        let q22 = self.get(x1, y1);

        let new_argb = {
            if self.is_height_step(q11, q12, q21, q22) {
                q11.argb
            } else {
                let argb_1 = linear_interp_argb(q11.argb(), q21.argb(), x0, x, x1);
                let argb_2 = linear_interp_argb(q12.argb(), q22.argb(), x0, x, x1);
                linear_interp_argb(argb_1, argb_2, y0, y, y1)
            }
        };

        self.data[(y * self.x_size) + x].argb = new_argb;
    }

    fn smoothen(&mut self, inflation_factor: usize) {
        for y in 0..self.y_size {
            for x in 0..self.x_size {
                let coarse_x0 = (x / inflation_factor) * inflation_factor;
                let coarse_y0 = (y / inflation_factor) * inflation_factor;
                let coarse_x1 = coarse_x0 + inflation_factor;
                let coarse_y1 = coarse_y0 + inflation_factor;

                if x != coarse_x0 || y != coarse_y0 {
                    self.bilinear_interp_height(x, y, coarse_x0, coarse_y0, coarse_x1, coarse_y1);
                    self.bilinear_interp_color(x, y, coarse_x0, coarse_y0, coarse_x1, coarse_y1);
                }
            }
        }
    }

    pub fn inflate(&mut self, factor: usize) {
        let mut new_data = vec![MapValue::new(); self.data.len() * factor * factor];
        for y in 0..self.y_size {
            for x in 0..self.x_size {
                let i = (y * self.x_size) + x;
                let j = ((y * factor) * (self.x_size * factor)) + (x * factor);

                new_data[j] = self.data[i].clone();
            }
        }
        self.data.resize(new_data.len(), MapValue::new());
        self.data.swap_with_slice(&mut new_data);
        self.x_size *= factor;
        self.y_size *= factor;

        self.smoothen(factor);
        debug_assert_eq!(self.data.len(), self.x_size * self.y_size);
    }

    pub fn add_color_lightness_noise(&mut self, amount: f64) {
        let amount = amount.clamp(0.0, 1.0);

        let count = self.x_size * self.y_size;
        let rng = ChaCha8Rng::seed_from_u64(0x736de2f2f99193e6);
        let noise: Vec<f64> = rng.sample_iter(Standard).take(count).collect();

        for y in 0..self.y_size {
            for x in 0..self.x_size {
                let i = (y * self.x_size) + x;

                let pnoise = (noise[i] - 0.5) * 200.0 * amount;

                let rgb = self.data[i].argb;
                let mut rgb = Rgb::new(
                    (rgb & 0xFF) as f64,
                    ((rgb >> 8) & 0xFF) as f64,
                    ((rgb >> 16) & 0xFF) as f64,
                    Some(((rgb >> 24) & 0xFF) as f64 / 255.0),
                );

                rgb.lighten(pnoise);

                self.data[i].argb = (rgb.blue().round() as i32).clamp(0, 0xFF) as u32
                    | (((rgb.green().round() as i32).clamp(0, 0xFF) as u32) << 8)
                    | (((rgb.red().round() as i32).clamp(0, 0xFF) as u32) << 16)
                    | ((((rgb.alpha() * 255.0).round() as i32).clamp(0, 0xFF) as u32) << 24);
            }
        }
    }
}

// vim: ts=4 sw=4 expandtab