mmd-mpl 0.2.10

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 {
            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 {
        format!(
            "{} {} {} {:.0};",
            self.bone, self.action, self.direction, self.amount
        )
    }

    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 {
        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![];
        }

        // 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 {
        format!(
            "@pose {} {{\n{}\n}}\n\nmain {{\n    {};\n}}",
            self.name,
            self.statements
                .iter()
                .map(|s| format!("    {}", s.to_string()))
                .collect::<Vec<String>>()
                .join("\n"),
            self.name
        )
    }

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

        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);
        }

        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 {
                    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(),
            ));
        }

        // Group statements by bone
        let mut bone_groups: std::collections::HashMap<String, Vec<MPLPoseStatement>> =
            std::collections::HashMap::new();
        for stmt in statements {
            bone_groups
                .entry(stmt.bone.clone())
                .or_insert_with(Vec::new)
                .push(stmt);
        }

        // Sort statements within each bone by action type (bend, turn, sway, move)
        let action_order = ["bend", "turn", "sway", "move"];
        for statements in bone_groups.values_mut() {
            statements.sort_by(|a, b| {
                let a_idx = action_order
                    .iter()
                    .position(|&x| x == a.action)
                    .unwrap_or(999);
                let b_idx = action_order
                    .iter()
                    .position(|&x| x == b.action)
                    .unwrap_or(999);
                a_idx.cmp(&b_idx)
            });
        }

        // Sort bones according to BONES array order and flatten
        let mut sorted_statements = Vec::new();
        for bone in crate::bone::BONES {
            if let Some(bone_statements) = bone_groups.get(*bone) {
                sorted_statements.extend(bone_statements.clone());
            }
        }

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