//                                         ______          _         
//                   _         _       _   | ___ \        | |        
//   _ __ ___   __ _| |_ _ __ | | ___ | |_ | |_/ /   _ ___| |_       
//  | '_ ` _ \ / _` | __| '_ \| |/ _ \| __||    / | | / __| __|      
//  | | | | | | (_| | |_| |_) | | (_) | |_ | |\ \ |_| \__ \ |_       
//  |_| |_| |_|\__,_|\__| .__/|_|\___/ \__|\_| \_\__,_|___/\__|      
//                      |_|                                          
extern crate tempfile;
extern crate image;
extern crate base64;

use std::fs::File;
use std::fs;
use std::io::{Write, Read, Seek, SeekFrom};
use std::process::Command;
use tempfile::NamedTempFile;
use std::fmt;
use image::GenericImage;
use image::Pixels;
use image::Pixel;
use base64::encode;

#[cfg(test)]
#[macro_use]
extern crate probability;

pub enum Markers {
    Point,
    Pixel,
    Circle,
    Triangle_Down,
    Triangle_Up,
    Triangle_Left,
    Triangle_Right,
    Tri_Down,
    Tri_Up,
    Tri_Left,
    Tri_Right,
    Square,
    Pentagon,
    Star,
    Hexagon1,
    Hexagon2,
    Plus,
    X ,
    Diamond,
    Thin_Diamond,
    VLine,
    HLine,
}

trait AsString {
    fn as_str(&self) -> &str;
}

impl Markers {
    pub fn as_str(&self) -> &str {
        match self {
            &Markers::Point => ".", 	// point marker
            &Markers::Pixel => ",", 	// pixel marker
            &Markers::Circle => "o", 	// circle marker
            &Markers::Triangle_Down => "v",  // triangle_down marker
            &Markers::Triangle_Up => "^", 	// triangle_up marker
            &Markers::Triangle_Left => "<",  //triangle_left marker
            &Markers::Triangle_Right => ">", // triangle_right marker
            &Markers::Tri_Down => "1",  // tri_down marker
            &Markers::Tri_Up => "2", // tri_up marker
            &Markers::Tri_Left => "3", // tri_left marker
            &Markers::Tri_Right => "4", // tri_right marker
            &Markers::Square => "s", // square marker
            &Markers::Pentagon => "p", // pentagon marker
            &Markers::Star => "*", // star marker
            &Markers::Hexagon1 => "h", // hexagon1 marker
            &Markers::Hexagon2 => "H", // hexagon2 marker
            &Markers::Plus => "+", // plus marker
            &Markers::X => "x", // x marker
            &Markers::Diamond => "D", // diamond marker
            &Markers::Thin_Diamond => "d", // thin_diamond marker
            &Markers::VLine => "|", // vline marker
            &Markers::HLine => "_", // hline marker
        }
    }
}

pub enum LineStyle {
    Dot,
    DashDot,
    Dash,
    Fill,
}

impl LineStyle {
    pub fn as_str(&self) -> &str {
        match self {
            &LineStyle::Dot => ":",
            &LineStyle::DashDot => "-.",
            &LineStyle::Dash => "--",
            &LineStyle::Fill => "-",
        }
    }
}

pub struct Figure {
    script: String
}

impl Figure {

    pub fn new() -> Figure {
        return Figure {
            script: "import matplotlib\nmatplotlib.use('agg')\nimport matplotlib.pyplot as plt\nfrom matplotlib.lines import Line2D\n\n".to_string()
        }   
    }

    pub fn add_plot(&mut self, p: String) {
        self.script += &p;
    }

    pub fn save(&mut self, output: &str, path: Option<&str>) -> (String, String, String) {
        self.script += &format!("plt.savefig('{}')\n", output);
        // create a temporary file
        let mut tmpfile: NamedTempFile = tempfile::NamedTempFile::new().unwrap();
        tmpfile.write_all(self.script.as_bytes());

        let python_path = match path {
            Some(s) => s,
            None => "/usr/local/bin/python3"
        };

        fs::metadata(python_path).expect("python binary not found at /usr/local/bin/python3");
        let mut echo_hello = Command::new(python_path);
        echo_hello.arg(tmpfile.path());
        echo_hello.output().expect("failed to execute process");

        return (self.script.to_string(), output.to_string(), tmpfile.path().to_str().unwrap().to_string());
    }

    pub fn as_base64(&mut self, path: Option<&str>) {
        // create a temporary file for storing the image
        let mut tmpfile: NamedTempFile = tempfile::NamedTempFile::new().unwrap();
        let filename = tmpfile.path().to_str().unwrap().to_string() + &".png".to_string();
        println!("tmpfile = {:?}", filename);
        self.save(&filename, path);
        // read the file back as a Base-64 PNG
        let mut buffer = Vec::new();
        let mut file = File::open(&filename).unwrap();
        let _out = file.read_to_end(&mut buffer);
        println!("EVCXR_BEGIN_CONTENT image/png\n{}\nEVCXR_END_CONTENT", base64::encode(&buffer));
    }

}

pub struct LinePlotOptions {
    pub marker: Option<Markers>,
    pub lineStyle: Option<LineStyle>,
    pub colour: Option<String>
}

impl LinePlotOptions {
    pub fn new() -> LinePlotOptions {
        return LinePlotOptions {
            marker: None,
            lineStyle: None,
            colour: None
        }
    }
}

pub struct ScatterPlotOptions {
    pub marker: Option<Markers>,
    pub alpha: Option<f64>,
}

impl ScatterPlotOptions {
    pub fn new() -> ScatterPlotOptions {
        return ScatterPlotOptions {
            marker: None,
            alpha: None,
        }
    }
}

impl fmt::Display for LinePlotOptions {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let mut options : Vec<String> = Vec::new();
        match self.marker {
            Some(ref m) => {options.push(format!("marker='{}'", m.as_str()));}
            None => {}
        }
        match self.lineStyle {
            Some(ref v) => {options.push(format!("linestyle='{}'", v.as_str()));}
            None => {}
        }
        match self.colour {
            Some(ref v) => {options.push(format!("c='{}'", v.as_str()));}
            None => {}
        }
        fmt.write_str(&options.join(", "));
        Ok(())
    }
}

impl fmt::Display for ScatterPlotOptions {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let mut options : Vec<String> = Vec::new();
        match self.marker {
            Some(ref m) => {options.push(format!("marker='{}'", m.as_str()));}
            None => {}
        }
        match self.alpha {
            Some(ref m) => {options.push(format!("alpha={}", m.to_string()));}
            None => {}
        }
        fmt.write_str(&options.join(", "));
        Ok(())
    }
}

fn convert_list_str<T>(x: Vec<T>) -> String where T: ToString {
    let mut result:Vec<String> = Vec::new();
    for _x in x { result.push(_x.to_string()); }
    return "[".to_string() + &result.join(", ") + &"]".to_string();
}

pub fn line_plot<U, T>(x: Vec<U>, y: Vec<T>, options: Option<LinePlotOptions>) -> String where U: ToString, T: ToString {
    let xs = convert_list_str::<U>(x);
    let ys = convert_list_str::<T>(y);
    match options {
        Some(opt) => {
            return format!("plt.plot({},{},{})\n", xs, ys, opt);
        },
        None => {
            return format!("plt.plot({},{})\n", xs, ys);
        }
    }
}

pub fn scatter_plot<U, T>(x: Vec<U>, y: Vec<T>, options: Option<ScatterPlotOptions>) -> String where U: ToString, T: ToString {
    let xs = convert_list_str::<U>(x);
    let ys = convert_list_str::<T>(y);
        match options {
        Some(opt) => {
            return format!("plt.scatter({},{},{})\n", xs, ys, opt);
        },
        None => {
            return format!("plt.scatter({},{})\n", xs, ys);
        }
    }
    
}

pub fn histogram<U>(x: Vec<U>, bins: Option<u32>) -> String where U: ToString {
    let xs = convert_list_str::<U>(x);
    return format!("plt.hist({}, bins={})\n", xs, bins.unwrap_or(20));
}

pub fn horizontal_line<U>(y: U, options: Option<LinePlotOptions>) -> String where U: ToString {
    match options {
        Some(opt) => {
            return format!("plt.axhline(y={},{})\n", y.to_string(), opt);
        },
        None => {
            return format!("plt.axhline(x={})\n", y.to_string());
        }
    }
}

pub fn vertical_line<U>(x: U, options: Option<LinePlotOptions>) -> String where U: ToString {
    match options {
        Some(opt) => {
            return format!("plt.axvline(x={},{})\n", x.to_string(), opt);
        },
        None => {
            return format!("plt.axvline(x={})\n", x.to_string());
        }
    }
}

pub fn line<U, T>(start : (U, T), end : (U, T), options : Option<LinePlotOptions>) -> String where U : ToString, T : ToString {
    match options {
        Some(opt) => {
            return format!("ax=plt.gca()\nax.add_line(Line2D([{},{}],[{},{}],{}))\n", start.0.to_string(), end.0.to_string(), start.1.to_string(), end.1.to_string(), opt);
        },
        None => {
            return format!("ax=plt.gca()\nax.add_line(Line2D([{},{}],[{},{}]))\n", start.0.to_string(), end.0.to_string(), start.1.to_string(), end.1.to_string());
        }
    }
}

#[cfg(test)]
mod lineplot {
    use super::*;
    use probability::distribution::Gaussian;
    use probability::sampler::Independent;
    use probability::source;
    use ::LineStyle::Dash;

    #[test]
    fn create_lineplot_basic() {
        let x = vec![1, 2, 3, 4];
        let y = vec![0.1, 0.2, 0.5, 0.3];
        let lp = line_plot::<i32, f64>(x, y, None);
        let mut figure = Figure::new();
        figure.add_plot(lp);
        figure.save("./examples/figures/lineplot_basic.png", None);
    }

     #[test]
    fn create_lineplot_basic_markers() {
        let x = vec![1, 2, 3, 4];
        let y = vec![0.1, 0.2, 0.5, 0.3];
        let mut options = LinePlotOptions::new();
        options.marker = Some(Markers::Diamond);
        let lp = line_plot::<i32, f64>(x, y, Some(options));
        let mut figure = Figure::new();
        figure.add_plot(lp);
        figure.save("./examples/figures/lineplot_basic_markers.png", None);
    }

     #[test]
    fn create_lineplot_basic_linestyle() {
        let x = vec![1, 2, 3, 4];
        let y = vec![0.1, 0.2, 0.5, 0.3];
        let mut options = LinePlotOptions::new();
        options.marker = Some(Markers::Diamond);
        options.lineStyle = Some(LineStyle::DashDot);
        let lp = line_plot::<i32, f64>(x, y, Some(options));
        let mut figure = Figure::new();
        figure.add_plot(lp);
        figure.save("./examples/figures/lineplot_basic_linestyle.png", None);
    }

    #[test]
    fn create_scatterplot_basic() {
        let x = vec![1, 2, 3, 4];
        let y = vec![0.1, 0.2, 0.5, 0.3];
        let lp = scatter_plot::<i32, f64>(x, y, None);
        let mut figure = Figure::new();
        figure.add_plot(lp);
        figure.save("./examples/figures/scatterplot_basic.png", None);
    }

    #[test]
    fn create_scatterplot_marker() {
        let x = vec![1, 2, 3, 4];
        let y = vec![0.1, 0.2, 0.5, 0.3];
        let mut options = ScatterPlotOptions::new();
        options.marker = Some(Markers::Diamond);
        let lp = scatter_plot::<i32, f64>(x, y, Some(options));
        let mut figure = Figure::new();
        figure.add_plot(lp);
        figure.save("./examples/figures/scatterplot_marker.png", None);
    }
    #[test]
    fn create_scatterplot_marker_alpha() {
        let x = vec![1, 2, 3, 4];
        let y = vec![0.1, 0.2, 0.5, 0.3];
        let mut options = ScatterPlotOptions::new();
        options.marker = Some(Markers::Diamond);
        options.alpha = Some(0.1);
        let lp = scatter_plot::<i32, f64>(x, y, Some(options));
        let mut figure = Figure::new();
        figure.add_plot(lp);
        figure.save("./examples/figures/scatterplot_marker_alpha.png", None);
    }

    #[test]
    fn create_histogram_default_bins() {
        let mut source = source::default();
        let gaussian = Gaussian::new(0.0, 2.0);
        let mut sampler = Independent(&gaussian, &mut source);
        let x = sampler.take(500).collect::<Vec<_>>();
        let plot = histogram::<f64>(x, None);
        let mut figure = Figure::new();
        figure.add_plot(plot);
        figure.save("./examples/figures/histogram_default_bins.png", None);
    }

    #[test]
    fn create_histogram_custom_bins() {
        let mut source = source::default();
        let gaussian = Gaussian::new(0.0, 2.0);
        let mut sampler = Independent(&gaussian, &mut source);
        let x = sampler.take(500).collect::<Vec<_>>();
        let plot = histogram::<f64>(x, Some(100));
        let mut figure = Figure::new();
        figure.add_plot(plot);
        figure.save("./examples/figures/histogram_custom_bins.png", None);
    }

    #[test]
    fn create_histogram_to_b64() {
        let mut source = source::default();
        let gaussian = Gaussian::new(0.0, 2.0);
        let mut sampler = Independent(&gaussian, &mut source);
        let x = sampler.take(500).collect::<Vec<_>>();
        let plot = histogram::<f64>(x, Some(100));
        let mut figure = Figure::new();
        figure.add_plot(plot);
        figure.as_base64(None);
    }

    fn mean(numbers: &Vec<f64>) -> f64 {

        let sum: f64 = numbers.iter().sum();

        sum / numbers.len() as f64

    }

    #[test]
    fn create_histogram_vertical_mean() {
        let mut source = source::default();
        let gaussian = Gaussian::new(0.0, 2.0);
        let mut sampler = Independent(&gaussian, &mut source);
        let x = sampler.take(500).collect::<Vec<_>>();
        let plot = histogram::<f64>(x.clone(), Some(100));
        let mut figure = Figure::new();
        figure.add_plot(plot);
        let mean = mean(&(x.clone()));
        let mut mean_line_opts = LinePlotOptions::new();
        mean_line_opts.lineStyle = Some(Dash);
        mean_line_opts.colour = Some("black".to_string());
        let mean_line = vertical_line(mean, Some(mean_line_opts));
        figure.add_plot(mean_line);
        figure.save("./examples/figures/histogram_mean_line.png", None);
    }

    #[test]
    fn create_horizontal_line() {
        let mut source = source::default();
        let gaussian = Gaussian::new(0.0, 2.0);
        let mut sampler = Independent(&gaussian, &mut source);
        let x = (0..500).collect();
        let y = sampler.take(500).collect::<Vec<_>>();
        let plot = scatter_plot::<i32, f64>(x, y.clone(), None);
        let mut figure = Figure::new();
        figure.add_plot(plot);
        let mean = mean(&(y.clone()));
        let mut mean_line_opts = LinePlotOptions::new();
        mean_line_opts.lineStyle = Some(Dash);
        mean_line_opts.colour = Some("black".to_string());
        let mean_line = horizontal_line(mean, Some(mean_line_opts));
        figure.add_plot(mean_line);
        print!("{:?}", figure.save("./examples/figures/horizontal_line.png", None));
    }

    #[test]
    fn line_segment() {
        let x = vec![1, 2, 3, 4];
        let y = vec![0.1, 0.2, 0.5, 0.3];
        let mut options = ScatterPlotOptions::new();
        options.marker = Some(Markers::Diamond);
        let lp = scatter_plot::<i32, f64>(x, y, Some(options));
        let mut figure = Figure::new();
        figure.add_plot(lp);
        let mut line_opt = LinePlotOptions::new();
        line_opt.colour = Some("red".to_string());
        line_opt.lineStyle = Some(LineStyle::Dash);
        figure.add_plot(line::<i32, f64>((1, 0.1), (4, 0.3), Some(line_opt)));
        figure.save("./examples/figures/line_segment.png", None);
        print!("{:?}", figure.script);

    }

}