//! Least square deforming of a 2D grid.

use types::Rectangle;
use {Line, Graphics, DrawState};
use triangulation::{tx, ty};
use math::{Matrix2d, Scalar, Vec2d};

/// Represents a deformed grid.
#[derive(Clone)]
pub struct DeformGrid {
    /// The number of columns in the grid.
    pub cols: usize,
    /// The number of rows in the grid.
    pub rows: usize,
    /// The grid undeformed, which is a plain rectangle.
    pub rect: [Scalar; 4],
    /// The vertices, deformed.
    pub vertices: Vec<Vec2d>,
    /// The triangle indices.
    pub indices: Vec<usize>,
    /// The texture coordinates.
    pub texture_coords: Vec<[f32; 2]>,
    /// Initial position of control points.
    pub ps: Vec<[Scalar; 2]>,
    /// The current position of control points.
    pub qs: Vec<[Scalar; 2]>,
    /// A weight computation buffer, one for each control point.
    pub wis: Vec<Scalar>,
}

impl DeformGrid {
    /// Creates a new DeformGrid.
    pub fn new(rect: Rectangle, cols: usize, rows: usize) -> DeformGrid {
        let x = rect[0];
        let y = rect[1];
        let w = rect[2];
        let h = rect[3];
        let mut vertices = Vec::new();
        let mut texture_coords: Vec<[f32; 2]> = Vec::new();
        let units_h = w / cols as Scalar;
        let units_v = h / rows as Scalar;
        let nx = cols + 1;
        let ny = rows + 1;
        for iy in 0..ny {
            for ix in 0..nx {
                vertices.push([x + ix as Scalar * units_h, y + iy as Scalar * units_v]);
                texture_coords.push([ix as f32 * units_h as f32 / w as f32,
                                     iy as f32 * units_v as f32 / h as f32]);
            }
        }

        let mut indices = Vec::new();
        for iy in 0..ny - 1 {
            for ix in 0..nx - 1 {
                indices.push(ix + iy * nx);
                indices.push((ix + 1) + iy * nx);
                indices.push(ix + (iy + 1) * nx);

                indices.push(ix + (iy + 1) * nx);
                indices.push((ix + 1) + iy * nx);
                indices.push((ix + 1) + (iy + 1) * nx);
            }
        }

        DeformGrid {
            cols: cols,
            rows: rows,
            rect: rect,
            vertices: vertices,
            indices: indices,
            texture_coords: texture_coords,
            ps: Vec::new(),
            qs: Vec::new(),
            wis: Vec::new(),
        }
    }

    /// Sets current control position.
    #[inline(always)]
    pub fn set_current(&mut self, i: usize, pos: Vec2d) {
        self.qs[i] = pos;
    }

    /// Sets original control position.
    #[inline(always)]
    pub fn set_original(&mut self, i: usize, pos: Vec2d) {
        self.ps[i] = pos;
    }

    /// Removes all control points.
    pub fn reset_control_points(&mut self) {
        // These values don't need drop, so we can set length directly.
        unsafe {
            self.ps.set_len(0);
            self.qs.set_len(0);
            self.wis.set_len(0);
        }
    }

    /// Sets vertices and texture coords back to default.
    pub fn reset_vertices_and_texture_coords(&mut self) {
        unsafe {
            self.vertices.set_len(0);
            self.texture_coords.set_len(0);
        }

        let cols = self.cols;
        let rows = self.rows;
        let r = self.rect;
        let units_h = r[2] / cols as Scalar;
        let units_v = r[3] / rows as Scalar;
        let nx = cols + 1;
        let ny = rows + 1;
        for iy in 0..ny {
            for ix in 0..nx {
                self.vertices.push([r[0] + ix as Scalar * units_h, r[1] + iy as Scalar * units_v]);
                self.texture_coords.push([ix as f32 * units_h as f32 / r[2] as f32,
                                          iy as f32 * units_v as f32 / r[3] as f32]);
            }
        }
    }

    /// Finds original coordinate.
    /// If the deformed grid is overlapping itself, multiple hits might occur.
    /// Returns the first hit it finds.
    pub fn hit(&self, pos: Vec2d) -> Option<Vec2d> {
        use math::{inside_triangle, to_barycentric, from_barycentric};
        let nx = self.cols + 1;
        let ny = self.rows + 1;
        for i in 0..nx - 1 {
            for j in 0..ny - 1 {
                let ip = i + j * nx;
                let p1 = self.vertices[ip];
                let ip = (i + 1) + j * nx;
                let p2 = self.vertices[ip];
                let ip = i + (j + 1) * nx;
                let p3 = self.vertices[ip];
                let ip = (i + 1) + (j + 1) * nx;
                let p4 = self.vertices[ip];
                let tri1 = [p1, p2, p3];
                let tri2 = [p3, p2, p4];
                if inside_triangle(tri1, [pos[0], pos[1]]) {
                    let b = to_barycentric(tri1, pos);
                    // Upper left triangle.
                    let tri = [[i as Scalar, j as Scalar],
                               [(i + 1) as Scalar, j as Scalar],
                               [i as Scalar, (j + 1) as Scalar]];
                    let tri_pos = from_barycentric(tri, b);
                    let r = self.rect;
                    let units_h = r[2] / self.cols as Scalar;
                    let units_v = r[3] / self.rows as Scalar;
                    return Some([r[0] + tri_pos[0] * units_h, r[1] + tri_pos[1] * units_v]);
                } else if inside_triangle(tri2, [pos[0], pos[1]]) {
                    let b = to_barycentric(tri2, pos);
                    // Lower right triangle.
                    let tri = [[i as Scalar, (j + 1) as Scalar],
                               [(i + 1) as Scalar, j as Scalar],
                               [(i + 1) as Scalar, (j + 1) as Scalar]];
                    let tri_pos = from_barycentric(tri, b);
                    let r = self.rect;
                    let units_h = r[2] / self.cols as Scalar;
                    let units_v = r[3] / self.rows as Scalar;
                    return Some([r[0] + tri_pos[0] * units_h, r[1] + tri_pos[1] * units_v]);
                }
            }
        }
        None
    }

    /// Draws deformed image using default method.
    #[inline(always)]
    pub fn draw_image<G>(&self,
                         texture: &<G as Graphics>::Texture,
                         draw_state: &DrawState,
                         transform: Matrix2d,
                         g: &mut G)
        where G: Graphics
    {
        g.deform_image(self, texture, draw_state, transform);
    }

    /// Draws deformed image using triangulation.
    pub fn draw_image_tri<G>(&self,
                             texture: &<G as Graphics>::Texture,
                             draw_state: &DrawState,
                             transform: Matrix2d,
                             g: &mut G)
        where G: Graphics
    {
        use BACK_END_MAX_VERTEX_COUNT as BUFFER_SIZE;

        let mat = transform;
        let color = [1.0; 4];
        let a = color[3];
        if a == 0.0 {
            return;
        }
        let mut vertices: [[f32; 2]; BUFFER_SIZE] = [[0.0; 2]; BUFFER_SIZE];
        let mut uvs: [[f32; 2]; BUFFER_SIZE] = [[0.0; 2]; BUFFER_SIZE];
        let mut offset = 0;
        for &ind in self.indices.iter() {
            if offset >= BUFFER_SIZE {
                g.tri_list_uv(&draw_state, &color, texture, |f| f(&vertices, &uvs));
                offset = 0;
            }
            let vert = self.vertices[ind];
            vertices[offset] = [tx(mat, vert[0], vert[1]), ty(mat, vert[0], vert[1])];
            uvs[offset] = self.texture_coords[ind];
            offset += 1;
        }
        if offset > 0 {
            g.tri_list_uv(&draw_state,
                          &color,
                          texture,
                          |f| f(&vertices[..offset], &uvs[..offset]));
        }
    }

    /// Adds a control point, in original coordinates.
    pub fn add_control_point(&mut self, pos: Vec2d) {
        self.ps.push(pos);
        self.qs.push(pos);
        self.wis.push(0.0);
    }

    /// Draw vertical grid lines.
    pub fn draw_vertical_lines<G>(&self,
                                  line: &Line,
                                  draw_state: &DrawState,
                                  transform: Matrix2d,
                                  g: &mut G)
        where G: Graphics
    {
        let grid = self;
        let nx = grid.cols + 1;
        let ny = grid.rows + 1;
        for i in 0..nx {
            for j in 0..ny - 1 {
                let ip = i + j * nx;
                let x1 = grid.vertices[ip][0];
                let y1 = grid.vertices[ip][1];
                let ip = i + (j + 1) * nx;
                let x2 = grid.vertices[ip][0];
                let y2 = grid.vertices[ip][1];
                line.draw([x1, y1, x2, y2], draw_state, transform, g);
            }
        }
    }

    /// Draw horizontal grid lines.
    pub fn draw_horizontal_lines<G>(&self,
                                    line: &Line,
                                    draw_state: &DrawState,
                                    transform: Matrix2d,
                                    g: &mut G)
        where G: Graphics
    {
        let grid = self;
        let nx = grid.cols + 1;
        let ny = grid.rows + 1;
        for i in 0..nx - 1 {
            for j in 0..ny {
                let ip = i + j * nx;
                let x1 = grid.vertices[ip][0];
                let y1 = grid.vertices[ip][1];
                let ip = (i + 1) + j * nx;
                let x2 = grid.vertices[ip][0];
                let y2 = grid.vertices[ip][1];
                line.draw([x1, y1, x2, y2], draw_state, transform, g);
            }
        }
    }

    /// Updates the grid, by deforming the vertices.
    pub fn update(&mut self) {
        use math::{add, cross, dot, mul_scalar, perp, square_len, sub};

        let &mut DeformGrid { cols,
                              rows,
                              rect,
                              ref mut ps,
                              ref mut qs,
                              ref mut vertices,
                              ref mut wis,
                              .. } = self;
        let ps = &mut ps[..];
        let qs = &mut qs[..];
        let wis = &mut wis[..];
        let vertices = &mut vertices[..];
        let x = rect[0];
        let y = rect[1];
        let w = rect[2];
        let h = rect[3];
        let units_h = w / cols as Scalar;
        let units_v = h / rows as Scalar;
        let num = ps.len();
        let eps = 0.00001;
        let nx = cols + 1;
        let ny = rows + 1;

        match ps.len() {
            0 => {
                return;
            }
            1 => {
                // Move all vertices same distance.
                let d = sub(qs[0], ps[0]);
                for iy in 0..ny {
                    for ix in 0..nx {
                        let ip = ix + iy * nx;
                        vertices[ip] = [x + ix as Scalar * units_h + d[0],
                                        y + iy as Scalar * units_v + d[1]];
                    }
                }
                return;
            }
            _ => {}
        }

        let zero = [0.0, 0.0];
        for m in 0..nx {
            for n in 0..ny {
                let ip = m + n * nx;
                let v = [m as Scalar * units_h + x, n as Scalar * units_v + y];
                let mut sum_wi = 0.0;
                let mut p_star = zero;
                let mut q_star = zero;
                for i in 0..num {
                    let pi = ps[i];
                    let vl = square_len(sub(pi, v));
                    let w = if vl < eps && vl > -eps {
                        1.0 / eps
                    } else {
                        1.0 / vl
                    };
                    sum_wi += w;
                    p_star = add(p_star, mul_scalar(pi, w));
                    q_star = add(q_star, mul_scalar(qs[i], w));
                    wis[i] = w;
                }

                let inv_sum_wi = 1.0 / sum_wi;
                p_star = mul_scalar(p_star, inv_sum_wi);
                q_star = mul_scalar(q_star, inv_sum_wi);
                let mut fr = zero;
                let vp = perp(sub(v, p_star));
                for i in 0..num {
                    let pi = ps[i];
                    let qi = qs[i];
                    let pi_hat = sub(pi, p_star);
                    let qi_hat = sub(qi, q_star);
                    let ai11 = cross(pi, vp);
                    let ai21 = dot(pi_hat, vp);
                    let ai12 = -dot(pi, vp);
                    let ai22 = cross(pi_hat, vp);
                    fr[0] += wis[i] * (qi_hat[0] * ai11 + qi_hat[1] * ai21);
                    fr[1] += wis[i] * (qi_hat[0] * ai12 + qi_hat[1] * ai22);
                }

                let vl = square_len(vp);
                let fl = square_len(fr);
                let vl = if fl == 0.0 { 0.0 } else { (vl / fl).sqrt() };
                vertices[ip] = add(mul_scalar(fr, vl), q_star);
            }
        }
    }
}