use std::fs;
use std::error::Error;
use std::time::Instant;

use rten_tensor::{NdTensor, AsView};
use rten::Model;
use ocrs::{OcrEngine, OcrEngineParams};

pub fn print_words(image: NdTensor<f32, 3>) -> Result<(), Box<dyn Error>> {

    let begin = Instant::now();

    println!("{:#?}", begin.elapsed());
    let begin = Instant::now();
    ocr(image).unwrap();
    println!("{:#?}", begin.elapsed());

    Ok(())
}

fn ocr(data: NdTensor<f32, 3>) -> Result<(), Box<dyn Error>> {

    let detection_model_data = fs::read("text-detection.rten")?;
    let rec_model_data = fs::read("text-recognition.rten")?;

    let detection_model = Model::load(&detection_model_data)?;
    let rec_model = Model::load(&rec_model_data)?;

    let ocr_engine = OcrEngine::new(OcrEngineParams {
        detection_model: Some(detection_model),
        recognition_model: Some(rec_model),
        ..Default::default()
    })?;

    let input = ocr_engine.prepare_input(data.view())?;

    let word_rects = ocr_engine.detect_words(&input)?;

    let line_rects = ocr_engine.find_text_lines(&input, &word_rects);

    let line_texts = ocr_engine.recognize_text(&input, &line_rects)?;

    for line in line_texts
        .iter()
        .flatten()
        .filter(|l| l.to_string().len() > 1) 
    {
        println!("{}", line);
    }

    

    Ok(())

    
}