use std::path::Path;
use image::{DynamicImage, GenericImageView, Luma, GrayImage};
use imageproc::{rect::Rect, point::Point};
use ndarray::{Array, ArrayBase, Dim, OwnedRepr};
use ort::{inputs, Session};

use crate::{error::PaddleOcrResult, PaddleOcrError};

pub struct Det {
    model: Session,
    rect_border_size: u32
}

impl Det {
    const RECT_BORDER_SIZE: u32 = 8;

    pub fn new(model: Session) -> Self {
        Self { model, rect_border_size: Self::RECT_BORDER_SIZE}
    }

    pub fn from_file(model_path: impl AsRef<Path>) -> PaddleOcrResult<Self> {
        let model = ort::SessionBuilder::new()?.commit_from_file(model_path)?;
        Ok(Self { model, rect_border_size: Self::RECT_BORDER_SIZE })
    }

    pub fn with_rect_border_size(mut self, rect_border_size: u32) -> Self {
        self.rect_border_size = rect_border_size;
        self
    }

    pub fn find_text_rect(&self, img: &DynamicImage) -> PaddleOcrResult<Vec<Rect>> {
        let input = Self::preprocess(img)?;
        let output = self.run_model(&input, img.width(), img.height())?;
        Ok(self.find_box(&output))
    }

    pub fn find_text_img(&self, img: &DynamicImage) -> PaddleOcrResult<Vec<DynamicImage>> {
        Ok(self.find_text_rect(img)?
            .iter()
            .map(|r| img.crop_imm(r.left() as u32, r.top() as u32, r.width(), r.height()))
            .collect()
        )
    }

    fn preprocess(img: &DynamicImage) -> PaddleOcrResult<ArrayBase<OwnedRepr<f32>, Dim<[usize; 4]>>> {
        let (w, h) = img.dimensions();
        let pad_w = Self::get_pad_length(w);
        let pad_h = Self::get_pad_length(h);

        let mut input = Array::zeros((1, 3, pad_h as usize, pad_w as usize));
        for pixel in img.pixels() {
            let x = pixel.0 as _;
            let y = pixel.1 as _;
            let [r, g, b, _] = pixel.2 .0;
            input[[0, 0, y, x]] = (((r as f32) / 255.) - 0.485) / 0.229;
            input[[0, 1, y, x]] = (((g as f32) / 255.) - 0.456) / 0.224;
            input[[0, 2, y, x]] = (((b as f32) / 255.) - 0.406) / 0.225;
        }
        Ok(input)
    }

    fn run_model(&self, input: &ArrayBase<OwnedRepr<f32>, Dim<[usize; 4]>>, width: u32, height: u32) -> PaddleOcrResult<GrayImage>{
        let pad_h = Self::get_pad_length(height);
        let outputs = self.model.run(inputs!["x" => input.view()]?)?;
        let output = outputs.iter().next().ok_or(PaddleOcrError::custom("no output"))?.1;
        let output = output
            .try_extract_tensor::<f32>()?
            .view()
            .t()
            .to_owned();
        let output: Vec<_> = output.iter().collect();
        let img = image::ImageBuffer::from_fn(width, height, |x, y| {
            Luma([(*output[(x * pad_h + y) as usize] * 255.0).min(255.0) as u8])
        });
        Ok(img)
    }

    fn find_box(&self, img: &GrayImage) ->Vec<Rect> {
        let (w, h) = img.dimensions();
        imageproc::contours::find_contours_with_threshold::<u32>(img, 200)
            .into_iter()
            .filter(|x| x.parent.is_none())
            .map(|x| x.points)
            .map(|x| Self::bounding_rect(&x))
            .filter_map(|x| x)
            .map(|x| {
                Rect::at((x.left() - self.rect_border_size as i32).max(0), (x.top() - self.rect_border_size as i32).max(0))
                    .of_size((x.width() + self.rect_border_size * 2).min(w), (x.height() + self.rect_border_size * 2).min(h))
            })
            .collect()
    }

    fn bounding_rect(points: &[Point<u32>]) -> Option<Rect> {
        let (x_min, x_max, y_min, y_max) = points.into_iter()
            .fold(None, |ret, p|{
                match ret {
                    None => Some((p.x, p.x, p.y, p.y)),
                    Some((x_min, x_max, y_min, y_max)) => {
                        Some((x_min.min(p.x), x_max.max(p.x), y_min.min(p.y), y_max.max(p.y)))
                    }
                }
            })?; 
        let width = (x_max - x_min) as u32;
        let height = (y_max - y_min) as u32;
        if width <= 5 || height <= 5 {
            return None;
        }
        Some(Rect::at(x_min as i32, y_min as i32).of_size(width, height))
    }

    const fn get_pad_length(length: u32) -> u32 {
        let i = length % 32;
        if i == 0 {
            length
        } else {
            length + 32 - i
        }
    }
}
