use camera::Camera;
use opencv::{core::{bitwise_and, bitwise_and_def, bitwise_or_def, in_range, KeyPoint, KeyPointTraitConst, Mat, MatTraitConst, MatTraitConstManual, Ptr, Size, VecN, Vector, CV_64F, CV_8UC4}, features2d::{draw_matches_def, BFMatcher, ORB}, gapi::bgr2_luv, highgui::{self, MouseCallback}, imgcodecs::ImreadModes, imgproc::{self, cvt_color_def, gaussian_blur, gaussian_blur_def, sobel_def, COLOR_BGR2GRAY, COLOR_BGR2HLS}, prelude::{DescriptorMatcherTrait, DescriptorMatcherTraitConst, Feature2DTrait}};
use webots::prelude::*;
use std::{f64::consts::TAU, time::Duration};
mod camera;

const TIME_STEP: Duration = Duration::from_millis(20);
const MAX_SPEED: f64 = TAU;

// const tweaker only supports floating point, these are mapped
#[const_tweaker::tweak]
const LOWERB_H: f32 = 18. / 180.;
#[const_tweaker::tweak]
const LOWERB_L: f32 = 23. / 255.;
#[const_tweaker::tweak]
const LOWERB_S: f32 = 36. / 255.;
#[const_tweaker::tweak]
const UPPERB_H: f32 = 92. / 180.;
#[const_tweaker::tweak]
const UPPERB_L: f32 = 151. / 255.;
#[const_tweaker::tweak]
const UPPERB_S: f32 = 169. / 255.;

struct MyRobot {
    camera: Camera,
    left_motor: Motor,
    right_motor: Motor,
    window: &'static str,
    orb_detector: Ptr<ORB>,
    matcher: BFMatcher,
    last_frame_descriptors: Option<Mat>,
    last_frame_keypoints: Option<Vector<KeyPoint>>,
    last_image: Option<Mat>,
    direction: f32,
    rotational_velocity: f32,
}

impl Robot for MyRobot {
    fn time_step(&self) -> Duration {
        TIME_STEP
    }

    fn init() -> Self {
        let camera = Camera::new("camera7");

        let left_motor = Motor::new("motorl");
        let right_motor = Motor::new("motorr");

        left_motor.set_position(f64::INFINITY);
        right_motor.set_position(f64::INFINITY);
        left_motor.set_velocity(0.1 * MAX_SPEED);
        right_motor.set_velocity(0.1 * MAX_SPEED);

        let window = "processed";
        //highgui::named_window(window, 1).unwrap();

        //highgui::set_mouse_callback(window, |x,y,_,_| {dbg!(x+y);}).unwrap();

        let orb_detector = ORB::create_def().unwrap();
        let matcher = opencv::features2d::BFMatcher::new_def().unwrap();
        //let last_frame_descriptors = opencv::core::Mat::default();

        Self {
            camera,
            left_motor,
            right_motor,
            window,
            orb_detector,
            matcher,
            last_frame_descriptors: None,
            last_frame_keypoints: None,
            last_image: None,
            direction: 0.,
            rotational_velocity: 0.,
        }
    }

    fn step(&mut self, _time: StepTime) {

        self.camera.enable(20);

        if let Ok(img) = self.camera.get_image() { // BGRA
            let img = opencv::core::Mat::new_rows_cols_with_bytes::<VecN<u8,4>>(240, 320, &img).unwrap();

            //opencv::imgcodecs::imwrite_def("cam.png", &img).unwrap();
            //highgui::imshow(self.window, &img).unwrap();


            let mut blurred = opencv::core::Mat::default();
            let mut edge = opencv::core::Mat::default();
            gaussian_blur_def(&img, &mut blurred, Size::new(3, 3), 0.).unwrap();
            //sobel_def(&blurred, &mut edge, CV_64F, 0, 2).unwrap();
            cvt_color_def(&blurred, &mut edge, COLOR_BGR2HLS).unwrap();
            highgui::imshow(self.window, &blurred).unwrap();
            let mut mask_a = opencv::core::Mat::default();
            let mut mask_b = opencv::core::Mat::default();
            let mut mask = opencv::core::Mat::default();
            opencv::imgcodecs::imwrite_def("blurred.png", &blurred).unwrap();
            println!("Lowrerh {}", LOWERB_H);
            let lowerb = dbg!([(*LOWERB_H*255.) as u8,(*LOWERB_L*255.) as u8,(*LOWERB_S*255.) as u8,255]);
            let upperb = [(*UPPERB_H*255.) as u8,(*UPPERB_L*255.) as u8,(*UPPERB_S*255.) as u8,255];
            in_range(&blurred, &lowerb, &upperb, &mut mask_a).unwrap(); // h low
                                                                                         // 103/180
            in_range(&blurred, &[0,30,4,255], &[52,209,63,255], &mut mask_b).unwrap(); // hh28
            highgui::imshow("a", &mask_a).unwrap();
            //highgui::imshow("b", &mask_b).unwrap();
            dbg!(blurred.row(240/2).unwrap().col(320/2).unwrap().data_bytes().unwrap());
            let mut range = opencv::core::Mat::default();
            bitwise_or_def(&mask_a, &mask_a, &mut mask).unwrap();
            bitwise_and(&img, &img, &mut range, &mask).unwrap();

            highgui::imshow("edge", &range).unwrap();

            let mask = opencv::core::Mat::default();
            let mut keypoints = opencv::core::Vector::default();
            let mut descriptors = opencv::core::Mat::default();
            self.orb_detector.detect_and_compute_def(&blurred, &mask, &mut keypoints, &mut descriptors).unwrap();
            let mut debugout = opencv::core::Mat::default();
            opencv::features2d::draw_keypoints(&img, &keypoints, &mut debugout, opencv::core::VecN([0.,255.,255.,255.]), opencv::features2d::DrawMatchesFlags::DEFAULT).unwrap();
            //highgui::imshow("keypoints", &debugout).unwrap();




            let mut matches: Vector<opencv::core::DMatch> = Vector::new();


            if let Some(mut last_frame_descriptors) = self.last_frame_descriptors.as_ref() {
                if let Err(e) = self.matcher.train_match_def(&mut last_frame_descriptors, &mut descriptors, &mut matches) {
                    eprintln!("no features to match");
                }


                let matches = matches.iter().filter(|m| m.distance < 105.).collect();
                draw_matches_def(self.last_image.as_ref().unwrap(), self.last_frame_keypoints.as_ref().unwrap(), &img , &keypoints, &matches, &mut debugout).unwrap();
                highgui::imshow("matches", &debugout).unwrap();

                let mut horiz = Vec::new();
                for hit in matches {
                    if hit.distance > 45. {continue;}
                    let last = self.last_frame_keypoints.as_ref().unwrap().get(hit.query_idx as usize).unwrap();
                    let current = keypoints.get(hit.train_idx as usize).unwrap();

                    let dx = last.pt().x - current.pt().x;

                    horiz.push(dx);
                }
                if horiz.len() > 0 {
                    let avg = horiz.iter().sum::<f32>()/horiz.len() as f32;

                    let lerp = 0.1 + horiz.len() as f32 * 0.05;
                    let lerp = lerp.max(0.65);
                    self.rotational_velocity = self.rotational_velocity * (1.-lerp) + avg * lerp;
                    //self.direction += avg;
                }
                self.direction += self.rotational_velocity;
            }
            self.last_frame_descriptors = Some(descriptors);
            self.last_frame_keypoints = Some(keypoints);
            self.last_image = Some(img.clone_pointee());
        }
        highgui::poll_key().unwrap();

        const DEG_TO_PX: f64 = 320. / 62.2;
        let goal = if _time.elapsed.as_secs() % 8 < 4 {
            90. * DEG_TO_PX
        } else {
            0. * DEG_TO_PX 
        };
        let error = goal - self.direction as f64;
        let mut error = error * 0.02;
        if error.abs() < 0.5 {
            error = 0.;
        }
        dbg!(self.direction);
        dbg!(self.rotational_velocity);
        let mut left_speed = error * MAX_SPEED;
        let mut right_speed = -error * MAX_SPEED;

        //self.left_motor.set_velocity(left_speed.clamp(-10., 10.));
        //self.right_motor.set_velocity(right_speed.clamp(-10., 10.));
        self.left_motor.set_velocity(0.);
        self.right_motor.set_velocity(0.);
    }
}

fn main() {
    println!("Rust controller has started");
    webots::run_robot::<MyRobot>();
}