mmd-mpl 0.3.4

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::{
    mpl::MPLBoneFrame,
    utils::{Quaternion, Vector3},
    with_bone_db, ActionRule,
};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MPLPoseStatement {
    pub bone: String,
    pub action: String,
    pub direction: String,
    pub amount: f32,
}

impl MPLPoseStatement {
    pub fn from_str(text: &str) -> Result<Self, String> {
        if text.is_empty() {
            return Err("Empty statement".to_string());
        }
        let parts = text.split_whitespace().collect::<Vec<&str>>();
        if parts.len() != 4 {
            if parts.len() == 2 && parts[1] == "reset" {
                return Ok(Self {
                    bone: parts[0].to_string(),
                    action: "reset".to_string(),
                    direction: "".to_string(),
                    amount: 0.0,
                });
            }
            return Err("Invalid statement".to_string());
        }

        let bone = parts[0].to_string();
        let action = parts[1].to_string();
        let direction = parts[2].to_string();
        let amount: f32 = parts[3]
            .trim()
            .parse()
            .map_err(|_| "Invalid degrees number".to_string())?;

        with_bone_db(|db| db.validate(&bone, &action, &direction, amount))?;

        Ok(Self {
            bone,
            action,
            direction,
            amount,
        })
    }

    pub fn to_string(&self) -> String {
        if self.action == "reset" {
            return format!("{} reset;", self.bone);
        }

        format!(
            "{} {} {} {:.0};",
            self.bone,
            self.action,
            self.direction,
            self.amount.round()
        )
    }

    pub fn to_vector(&self) -> Vector3 {
        let rule = with_bone_db(|db| {
            db.get_rule(&self.bone, &self.action, &self.direction)
                .cloned()
        });

        let rule = match rule {
            Some(r) => r,
            None => return Vector3::new(0.0, 0.0, 0.0),
        };

        let normalized_axis = rule.axis.normalize();
        normalized_axis.multiply_by_scalar(self.amount)
    }

    pub fn from_vector(bone: &str, target_vector: Vector3) -> Vec<Self> {
        let bone = bone.to_string();
        let mut statements = vec![];

        // Map vector components to move directions
        let direction_mappings = [
            (target_vector.x, "right", "left"),
            (target_vector.y, "up", "down"),
            (target_vector.z, "backward", "forward"),
        ];

        for (component, pos_dir, neg_dir) in direction_mappings {
            if component.abs() > 0.01 {
                let direction = if component > 0.0 { pos_dir } else { neg_dir };
                let amount = component.abs();

                // Check if this bone supports this move direction
                let has_rule = with_bone_db(|db| db.get_rule(&bone, "move", direction).is_some());
                if has_rule {
                    statements.push(Self {
                        bone: bone.clone(),
                        action: "move".to_string(),
                        direction: direction.to_string(),
                        amount,
                    });
                }
            }
        }
        statements
    }

    pub fn to_quaternion(&self) -> Quaternion {
        if self.action == "reset" {
            return Quaternion::identity();
        }

        let rule = with_bone_db(|db| {
            db.get_rule(&self.bone, &self.action, &self.direction)
                .cloned()
        });

        let rule = match rule {
            Some(r) => r,
            None => return Quaternion::identity(),
        };

        let normalized_axis = rule.axis.normalize();

        let radians = self.amount * (std::f32::consts::PI / 180.0);
        let half_angle = radians / 2.0;
        let sin = half_angle.sin();
        let cos = half_angle.cos();

        Quaternion::new(
            normalized_axis.x * sin,
            normalized_axis.y * sin,
            normalized_axis.z * sin,
            cos,
        )
    }

    pub fn from_quaternion(bone: &str, target_quat: Quaternion) -> Vec<Self> {
        let bone = bone.to_string();

        // Gather all possible (action, direction) rules for this bone
        let possible_actions: Vec<(String, String, ActionRule)> = with_bone_db(|db| {
            let mut vec = Vec::new();
            if let Some(actions) = db.actions(&bone) {
                for action in actions {
                    if action == "move" {
                        continue;
                    }
                    if let Some(directions) = db.directions(&bone, action) {
                        for direction in directions {
                            if let Some(rule) = db.get_rule(&bone, action, direction) {
                                vec.push((action.to_string(), direction.to_string(), rule.clone()));
                            }
                        }
                    }
                }
            }
            vec
        });
        if possible_actions.is_empty() {
            return vec![];
        }

        // Special case: if target quaternion is identity (0,0,0,1),
        // return bone reset
        if target_quat.x == 0.0
            && target_quat.y == 0.0
            && target_quat.z == 0.0
            && target_quat.w == 1.0
        {
            let mut statements = Vec::new();
            statements.push(Self {
                bone: bone.clone(),
                action: "reset".to_string(),
                direction: "".to_string(),
                amount: 0.0,
            });
            return statements;
        }

        // Ensure deterministic order independent of HashMap iteration
        let mut possible_actions = possible_actions;
        possible_actions.sort_by(|a, b| {
            let key_a = format!("{}-{}", a.0, a.1);
            let key_b = format!("{}-{}", b.0, b.1);
            key_a.cmp(&key_b)
        });

        // Evaluate fitness of a degree combination
        let evaluate_combination = |degrees: &[f32]| -> f32 {
            if degrees.len() != possible_actions.len() {
                return f32::INFINITY;
            }

            let mut combined_quaternion = Quaternion::identity();

            for (i, deg) in degrees.iter().enumerate() {
                let clamped_deg = deg.max(0.0).min(possible_actions[i].2.limit);
                if clamped_deg > 0.01 {
                    // Only apply significant rotations
                    let q = Quaternion::from_axis_angle(possible_actions[i].2.axis, clamped_deg);
                    combined_quaternion = combined_quaternion.multiply(&q);
                }
            }

            target_quat.angular_distance(&combined_quaternion)
        };

        // Nelder-Mead simplex optimization algorithm
        let nelder_mead = |initial_guess: &[f32], max_iterations: usize| -> (Vec<f32>, f32) {
            let n = initial_guess.len();
            let alpha = 1.0; // reflection coefficient
            let gamma = 2.0; // expansion coefficient
            let rho = 0.5; // contraction coefficient
            let sigma = 0.5; // shrinkage coefficient

            // Initialize simplex with n+1 points
            let mut simplex: Vec<(Vec<f32>, f32)> = Vec::new();

            simplex.push((initial_guess.to_vec(), evaluate_combination(initial_guess)));

            // Create additional points by perturbing initial guess
            for i in 0..n {
                let mut point = initial_guess.to_vec();
                let range = possible_actions[i].2.limit;
                point[i] += range * 0.1;
                let value = evaluate_combination(&point);
                simplex.push((point, value));
            }

            // Main optimization loop
            for _ in 0..max_iterations {
                simplex.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());

                let best_value = simplex[0].1;
                let worst_value = simplex[n].1;
                let second_worst_value = simplex[n - 1].1;

                // Check convergence
                if worst_value - best_value < 0.0001 {
                    break;
                }

                // Calculate centroid (excluding worst point)
                let mut centroid = vec![0.0f32; n];
                for i in 0..n {
                    for j in 0..n {
                        centroid[j] += simplex[i].0[j];
                    }
                }
                for j in 0..n {
                    centroid[j] /= n as f32;
                }

                // Reflection step
                let reflected: Vec<f32> = centroid
                    .iter()
                    .zip(&simplex[n].0)
                    .map(|(c, w)| c + alpha * (c - w))
                    .collect();
                let reflected_value = evaluate_combination(&reflected);

                if reflected_value >= best_value && reflected_value < second_worst_value {
                    simplex[n] = (reflected, reflected_value);
                    continue;
                }

                // Expansion step
                if reflected_value < best_value {
                    let expanded: Vec<f32> = centroid
                        .iter()
                        .zip(&reflected)
                        .map(|(c, r)| c + gamma * (r - c))
                        .collect();
                    let expanded_value = evaluate_combination(&expanded);

                    if expanded_value < reflected_value {
                        simplex[n] = (expanded, expanded_value);
                    } else {
                        simplex[n] = (reflected, reflected_value);
                    }
                    continue;
                }

                // Contraction step
                let contracted: Vec<f32> = centroid
                    .iter()
                    .zip(&simplex[n].0)
                    .map(|(c, w)| c + rho * (w - c))
                    .collect();
                let contracted_value = evaluate_combination(&contracted);

                if contracted_value < worst_value {
                    simplex[n] = (contracted, contracted_value);
                    continue;
                }

                // Shrinkage step
                let best_point = simplex[0].0.clone();
                for i in 1..=n {
                    for j in 0..n {
                        simplex[i].0[j] = best_point[j] + sigma * (simplex[i].0[j] - best_point[j]);
                    }
                    simplex[i].1 = evaluate_combination(&simplex[i].0);
                }
            }

            simplex.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
            (simplex[0].0.clone(), simplex[0].1)
        };

        let mut best_result = (Vec::new(), f32::INFINITY);

        // Try optimization from multiple starting points for global search
        let starting_points = vec![
            vec![0.0; possible_actions.len()], // Zero start
            possible_actions
                .iter()
                .enumerate()
                .map(|(i, action)| {
                    let limit = action.2.limit.min(30.0);
                    // Use deterministic "random" values based on index
                    let pseudo_random = ((i * 12345) % 1000) as f32 / 1000.0;
                    (limit * pseudo_random).min(limit)
                })
                .collect(), // Pseudo-random start
            possible_actions
                .iter()
                .map(|action| action.2.limit * 0.5)
                .collect(), // Mid-range start
            possible_actions
                .iter()
                .enumerate()
                .map(|(i, action)| {
                    if i % 2 == 0 {
                        action.2.limit * 0.3
                    } else {
                        action.2.limit * 0.7
                    }
                })
                .collect(), // Mixed start
        ];

        for start in starting_points {
            let result = nelder_mead(&start, 1000);
            if result.1 < best_result.1 {
                best_result = result;
            }
        }

        // Convert optimal degrees to MPL statements and simplify opposing actions
        let mut action_map: HashMap<String, HashMap<String, f32>> = HashMap::new();

        // Group degrees by action and direction
        for (i, deg) in best_result.0.iter().enumerate() {
            if *deg > 0.01 {
                let action = &possible_actions[i];
                let clamped_deg = deg.max(0.0).min(action.2.limit);

                action_map
                    .entry(action.0.clone())
                    .or_default()
                    .insert(action.1.clone(), clamped_deg);
            }
        }

        // Simplify opposing directions within each action
        let mut statements = Vec::new();
        for (action, directions) in action_map.into_iter() {
            // Handle opposing pairs
            let opposing_pairs = [("forward", "backward"), ("left", "right")];
            let mut processed_directions = std::collections::HashSet::new();

            for (dir1, dir2) in opposing_pairs.iter() {
                if directions.contains_key(*dir1)
                    && directions.contains_key(*dir2)
                    && !processed_directions.contains(*dir1)
                    && !processed_directions.contains(*dir2)
                {
                    let deg1 = directions.get(*dir1).unwrap();
                    let deg2 = directions.get(*dir2).unwrap();
                    let net_degrees = (deg1 - deg2).abs();

                    if net_degrees > 0.01 {
                        let net_direction = if deg1 > deg2 { dir1 } else { dir2 };
                        statements.push(Self {
                            bone: bone.clone(),
                            action: action.clone(),
                            direction: net_direction.to_string(),
                            amount: net_degrees,
                        });
                    }

                    processed_directions.insert(*dir1);
                    processed_directions.insert(*dir2);
                }
            }

            // Handle remaining directions that don't have opposing pairs
            for (direction, degrees) in directions.iter() {
                if !processed_directions.contains(direction.as_str()) && *degrees > 0.01 {
                    statements.push(Self {
                        bone: bone.clone(),
                        action: action.clone(),
                        direction: direction.clone(),
                        amount: *degrees,
                    });
                }
            }
        }

        // Format statements to match TypeScript output format
        let s = statements
            .into_iter()
            .map(|stmt| MPLPoseStatement {
                amount: (stmt.amount / 5.0).round() * 5.0,
                ..stmt
            })
            .filter(|stmt| stmt.amount.abs() >= 0.0)
            .collect();
        return s;
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MPLPose {
    pub name: String,
    pub statements: Vec<MPLPoseStatement>,
}

impl MPLPose {
    pub fn new(name: String, statements: Vec<MPLPoseStatement>) -> Self {
        Self { name, statements }
    }

    pub fn to_string(&self) -> String {
        // Group statements by bone
        let mut bone_groups: HashMap<String, Vec<&MPLPoseStatement>> = HashMap::new();
        for statement in &self.statements {
            bone_groups
                .entry(statement.bone.clone())
                .or_insert_with(Vec::new)
                .push(statement);
        }

        // Sort bones according to BONES array order for consistent output
        let mut sorted_bones: Vec<_> = bone_groups.into_iter().collect();
        sorted_bones.sort_by(|(a, _), (b, _)| {
            let a_idx = crate::bone::BONES
                .iter()
                .position(|&x| x == a)
                .unwrap_or(999);
            let b_idx = crate::bone::BONES
                .iter()
                .position(|&x| x == b)
                .unwrap_or(999);
            a_idx.cmp(&b_idx)
        });

        // Convert each bone's statements to compound format
        let mut compound_statements = Vec::new();
        for (bone, statements) in sorted_bones {
            if statements.len() == 1 {
                // Single statement - use as-is
                compound_statements.push(format!("    {}", statements[0].to_string()));
            } else {
                // Multiple statements - combine into compound format
                let compound_parts: Vec<String> = statements
                    .iter()
                    .map(|stmt| {
                        if stmt.action == "reset" {
                            "reset".to_string()
                        } else {
                            format!("{} {} {}", stmt.action, stmt.direction, stmt.amount.round())
                        }
                    })
                    .collect();
                compound_statements.push(format!("    {} {};", bone, compound_parts.join(", ")));
            }
        }

        format!("{{\n{}\n}}", compound_statements.join("\n"))
    }

    pub fn to_script(&self) -> String {
        // Filter out reset statements and zero-degree actions for static pose
        let filtered_statements: Vec<MPLPoseStatement> = self
            .statements
            .iter()
            .filter(|stmt| stmt.action != "reset" && stmt.amount.abs() > 0.01)
            .cloned()
            .collect();

        let filtered_pose = MPLPose::new(self.name.clone(), filtered_statements);

        format!(
            "@pose {} {}\n\nmain {{\n    {};\n}}",
            self.name,
            filtered_pose.to_string(),
            self.name
        )
    }

    pub fn to_block(&self) -> String {
        format!("@pose {} {}\n", self.name, self.to_string())
    }

    pub fn to_bone_frames(&self) -> Vec<MPLBoneFrame> {
        let mut frames = vec![];

        // Expand grouped finger statements into individual joint statements
        let mut expanded_statements = Vec::new();
        for statement in &self.statements {
            if Self::is_grouped_finger_bone(&statement.bone) {
                expanded_statements.extend(Self::expand_grouped_finger_statement(statement));
            } else {
                expanded_statements.push(statement.clone());
            }
        }

        // Process all statements (expanded and original)
        let mut bone_groups: HashMap<String, Vec<&MPLPoseStatement>> = HashMap::new();
        for statement in &expanded_statements {
            bone_groups
                .entry(statement.bone.clone())
                .or_insert_with(Vec::new)
                .push(statement);
        }

        for (bone, bone_statements) in bone_groups {
            let mut combined_position = Vector3::new(0.0, 0.0, 0.0);
            let mut combined_quaternion = Quaternion::identity();

            for statement in bone_statements {
                if statement.action == "move" {
                    let vector = statement.to_vector();
                    combined_position = combined_position.add(&vector);
                } else if statement.action == "reset" {
                    combined_quaternion = Quaternion::identity();
                } else {
                    let quaternion = statement.to_quaternion();
                    combined_quaternion = combined_quaternion.multiply(&quaternion);
                }
            }

            let bone_name_jp =
                with_bone_db(|db| db.japanese_name(&bone).unwrap_or(&bone).to_string());

            frames.push(MPLBoneFrame::new(
                bone,
                bone_name_jp,
                combined_position,
                combined_quaternion,
            ));
        }

        frames
    }

    pub fn from_bone_frames(name: &str, frames: Vec<MPLBoneFrame>) -> Self {
        let mut statements = vec![];
        for frame in frames.iter() {
            statements.extend(MPLPoseStatement::from_vector(
                &frame.name_en(),
                frame.position(),
            ));
            statements.extend(MPLPoseStatement::from_quaternion(
                &frame.name_en(),
                frame.rotation(),
            ));
        }

        // Consolidate individual finger joint statements into grouped finger statements
        let consolidated_statements = Self::consolidate_finger_statements(&statements);

        Self::new(name.to_string(), consolidated_statements)
    }

    fn expand_grouped_finger_statement(
        stmt: &crate::pose::MPLPoseStatement,
    ) -> Vec<crate::pose::MPLPoseStatement> {
        let mut expanded_statements = Vec::new();

        // Define finger joint mappings with ratios
        // Based on natural finger movement: MCP (1.0) > PIP (0.85-0.95) > DIP (0.5-0.7)
        let finger_mappings = match stmt.bone.as_str() {
            "thumb_l" => vec![("thumb_0_l", 1.0), ("thumb_1_l", 0.85)], // Thumb has different joint structure
            "index_l" => vec![("index_0_l", 1.0), ("index_1_l", 0.9), ("index_2_l", 0.65)],
            "middle_l" => vec![
                ("middle_0_l", 1.0),
                ("middle_1_l", 0.9),
                ("middle_2_l", 0.65),
            ],
            "ring_l" => vec![("ring_0_l", 1.0), ("ring_1_l", 0.88), ("ring_2_l", 0.6)], // Ring finger bends slightly less
            "pinky_l" => vec![("pinky_0_l", 1.0), ("pinky_1_l", 0.85), ("pinky_2_l", 0.55)], // Pinky bends the least
            "thumb_r" => vec![("thumb_1_r", 1.0), ("thumb_2_r", 0.85)],
            "index_r" => vec![("index_0_r", 1.0), ("index_1_r", 0.9), ("index_2_r", 0.65)],
            "middle_r" => vec![
                ("middle_0_r", 1.0),
                ("middle_1_r", 0.9),
                ("middle_2_r", 0.65),
            ],
            "ring_r" => vec![("ring_0_r", 1.0), ("ring_1_r", 0.88), ("ring_2_r", 0.6)],
            "pinky_r" => vec![("pinky_0_r", 1.0), ("pinky_1_r", 0.85), ("pinky_2_r", 0.55)],
            _ => return vec![stmt.clone()], // Not a grouped finger bone, return original
        };

        // Create expanded statements for each joint
        for (joint_bone, ratio) in finger_mappings {
            let adjusted_amount = (stmt.amount * ratio).round();
            if adjusted_amount.abs() >= 1.0 {
                // Only add if amount is significant
                expanded_statements.push(crate::pose::MPLPoseStatement {
                    bone: joint_bone.to_string(),
                    action: stmt.action.clone(),
                    direction: stmt.direction.clone(),
                    amount: adjusted_amount,
                });
            }
        }

        expanded_statements
    }

    fn is_grouped_finger_bone(bone: &str) -> bool {
        bone == "thumb_l"
            || bone == "index_l"
            || bone == "middle_l"
            || bone == "ring_l"
            || bone == "pinky_l"
            || bone == "thumb_r"
            || bone == "index_r"
            || bone == "middle_r"
            || bone == "ring_r"
            || bone == "pinky_r"
    }

    fn consolidate_finger_statements(statements: &[MPLPoseStatement]) -> Vec<MPLPoseStatement> {
        let mut consolidated = Vec::new();

        for stmt in statements {
            // Only keep _0 joints and rename them to grouped form, ignore all others
            if stmt.bone.ends_with("_0_l") || stmt.bone.ends_with("_0_r") {
                let grouped_bone = if stmt.bone.ends_with("_0_l") {
                    stmt.bone.replace("_0_l", "_l")
                } else {
                    stmt.bone.replace("_0_r", "_r")
                };
                consolidated.push(MPLPoseStatement {
                    bone: grouped_bone,
                    action: stmt.action.clone(),
                    direction: stmt.direction.clone(),
                    amount: stmt.amount,
                });
            } else if !stmt.bone.ends_with("_1_l")
                && !stmt.bone.ends_with("_2_l")
                && !stmt.bone.ends_with("_1_r")
                && !stmt.bone.ends_with("_2_r")
            {
                // Keep non-finger joint statements as they are
                consolidated.push(stmt.clone());
            }
            // Ignore all other individual joints (_1_l, _2_l, etc.)
        }

        consolidated
    }
}