use crate::scene::SceneAnimation;
use gloss_img::dynamic_image::DynImage;
use gloss_renderer::{
    components::{DiffuseImg, Faces, MetalnessImg, Name, NormalImg, RoughnessImg, UVs, Verts},
    scene::Scene,
};
use gloss_utils::{
    bshare::{ToNalgebraFloat, ToNalgebraInt, ToNdArray},
    nshare::ToNalgebra,
    tensor::DynamicMatrixOps,
};
use image::imageops::FilterType;
use log::info;
use nalgebra::DMatrix;
use ndarray::{self as nd, s};
use smpl_core::common::types::SmplType;
use smpl_core::{
    codec::gltf::GltfCodec,
    common::{metadata::smpl_metadata, pose::Pose, smpl_params::SmplParams},
    conversions::pose_remap::PoseRemap,
};
use smpl_core::{codec::gltf::PropData, common::transform_sequence::TransformSequence};
use smpl_core::{
    codec::{gltf::PerBodyData, scene::CameraTrack},
    common::{
        animation::Animation,
        betas::Betas,
        expression::Expression,
        pose_override::PoseOverride,
        pose_retarget::RetargetPoseYShift,
        smpl_model::SmplCache,
        smpl_options::SmplOptions,
        types::{FaceType, UpAxis},
    },
};
use smpl_utils::array::{Gather2D, Gather3D};
use std::f32::consts::PI;
use std::ops::Deref;
/// Creates a ``GltfCodec`` from an entity by extracting components from it
pub trait GltfCodecGloss {
    fn from_scene(scene: &Scene, options: &GltfInteropOptions) -> GltfCodec;
    fn from_entities(scene: &Scene, options: &GltfInteropOptions, entities: Vec<String>) -> GltfCodec;
    fn from_scene_with_body_indices(scene: &Scene, options: &GltfInteropOptions, body_idxs: Vec<usize>) -> GltfCodec;
}
fn get_image(image: &DynImage, to_gray: bool, max_texture_size: Option<u32>) -> DynImage {
    let mut image = image.clone();
    if to_gray {
        image = image.grayscale();
    }
    if let Some(force_image_size) = max_texture_size {
        if image.width() > force_image_size {
            image.resize(force_image_size, force_image_size, FilterType::Gaussian)
        } else {
            image
        }
    } else {
        image
    }
}
pub struct GltfInteropOptions {
    pub max_texture_size: Option<u32>,
    pub export_camera: bool,
    pub export_shape: bool,
}
impl Default for GltfInteropOptions {
    fn default() -> Self {
        Self {
            max_texture_size: None,
            export_camera: true,
            export_shape: true,
        }
    }
}
/// Trait implementation for ``GltfCodec``
impl GltfCodecGloss for GltfCodec {
    /// Get a ``GltfCodec`` from the scene
    fn from_scene(scene: &Scene, options: &GltfInteropOptions) -> GltfCodec {
        let smpl_models = scene.get_resource::<&SmplCache>().unwrap();
        gltfcodec_from_scene(scene, &smpl_models, options, None)
    }
    /// Get a ``GltfCodec`` from the scene
    fn from_entities(scene: &Scene, options: &GltfInteropOptions, entities: Vec<String>) -> GltfCodec {
        let smpl_models = scene.get_resource::<&SmplCache>().unwrap();
        gltfcodec_from_scene(scene, &smpl_models, options, Some(&entities))
    }
    /// Get a ``GltfCodec`` from the scene
    fn from_scene_with_body_indices(scene: &Scene, options: &GltfInteropOptions, body_idxs: Vec<usize>) -> GltfCodec {
        assert!(!body_idxs.is_empty(), "At least one index must be specified");
        let smpl_models = scene.get_resource::<&SmplCache>().unwrap();
        let num_smpl_bodies = scene.world.query::<&SmplParams>().iter().len();
        if *body_idxs.iter().max().unwrap() >= num_smpl_bodies {
            info!("Some of the indices are out of range. Ignoring these indices.");
        }
        let entities_to_export: Vec<String> = body_idxs.into_iter().map(|idx| format!("avatar-{:02}", idx + 1)).collect();
        gltfcodec_from_scene(scene, &smpl_models, options, Some(&entities_to_export))
    }
}
/// Function to get a ``GltfCodec`` from a scene, possible limited to certain entities
#[allow(clippy::too_many_lines)]
#[allow(clippy::trivially_copy_pass_by_ref)]
fn gltfcodec_from_scene(scene: &Scene, smpl_models: &SmplCache, options: &GltfInteropOptions, entities: Option<&Vec<String>>) -> GltfCodec {
    let now = wasm_timer::Instant::now();
    let mut gltf_codec = GltfCodec::default();
    let mut nr_frames = 0;
    if options.export_camera {
        let mut cameras_query = scene.world.query::<&CameraTrack>();
        for (_, camera_track) in cameras_query.iter() {
            gltf_codec.camera_track = Some(camera_track.clone());
        }
    }
    let mut query = scene.world.query::<(&SmplParams, &Name)>();
    let num_bodies = query.iter().len();
    gltf_codec.num_bodies = num_bodies;
    let mut should_export_posedirs = false;
    let mut should_export_exprdirs = false;

    let mut num_expression_blend_shapes = 0;
    for (entity, (smpl_params, _name)) in query.iter() {
        if scene.world.has::<Animation>(entity).unwrap() && smpl_params.enable_pose_corrective {
            should_export_posedirs = true;
        }
        let smpl_model = smpl_models.get_model_ref(smpl_params.smpl_type, smpl_params.gender).unwrap();
        if let Ok(anim) = scene.get_comp::<&Animation>(&entity) {
            if anim.has_expression() && smpl_model.get_expression_dirs().is_some() {
                should_export_exprdirs = true;
                num_expression_blend_shapes = smpl_model.get_expression_dirs().unwrap().shape().dims[1];
            }
        }
    }

    if let Ok(scene_anim) = scene.get_resource::<&SceneAnimation>() {
        nr_frames = scene_anim.num_frames;
        let mut keyframe_times: Vec<f32> = Vec::new();
        let fps = scene_anim.config.fps;
        #[allow(clippy::cast_precision_loss)]
        for global_frame_idx in 0..nr_frames {
            keyframe_times.push((global_frame_idx as f32) / fps);
        }
        gltf_codec.keyframe_times = Some(keyframe_times);
        gltf_codec.frame_count = Some(nr_frames);
    }
    for (body_idx, (entity, (smpl_params, name))) in query.iter().enumerate() {
        if let Some(entities) = entities {
            if !entities.contains(&name.0) {
                continue;
            }
        }
        let smpl_version = smpl_params.smpl_type;
        let gender = smpl_params.gender as i32;
        let mut current_body = PerBodyData::default();
        assert!(smpl_version != SmplType::SmplPP, "GLTF export for SMPL++ is not supported yet!");
        let smpl_model = smpl_models.get_model_ref(smpl_params.smpl_type, smpl_params.gender).unwrap();
        let betas = if options.export_shape {
            scene
                .get_comp::<&Betas>(&entity)
                .expect("Betas component does not exist!")
                .deref()
                .clone()
        } else {
            Betas::default()
        };
        let default_pose = Pose::new_empty(UpAxis::Y, smpl_params.smpl_type);
        let default_expression = Expression::new_empty(10, FaceType::SmplX);
        let mut smpl_output = smpl_model.forward(&SmplOptions::default(), &betas, &default_pose, Some(&default_expression));
        smpl_output.compute_normals();
        smpl_output = smpl_model.create_body_with_uv(&smpl_output);
        let metadata = smpl_metadata(&smpl_params.smpl_type);
        let mut num_total_blendshapes = 0;
        if should_export_posedirs {
            num_total_blendshapes += metadata.num_pose_blend_shapes + 1;
        }
        if should_export_exprdirs {
            num_total_blendshapes += num_expression_blend_shapes + 1;
        }
        gltf_codec.smpl_type = smpl_version;
        gltf_codec.gender = gender;
        current_body.joint_poses = Some(default_pose.joint_poses.to_ndarray().clone());
        gltf_codec.default_joint_poses = Some(default_pose.clone().joint_poses.to_ndarray());
        current_body.body_translation = Some(default_pose.clone().global_trans.to_ndarray().to_shape((1, 3)).unwrap().to_owned());
        let verts_na = smpl_output.verts.to_nalgebra();
        let normals_na = smpl_output.normals.as_ref().expect("SMPL Output is missing normals!").to_nalgebra();
        let faces_na = smpl_output.faces.to_nalgebra();
        let uvs_na = smpl_output.uvs.as_ref().expect("SMPL Output is missing UVs!").to_nalgebra();
        current_body.positions = Some(verts_na);
        current_body.normals = Some(normals_na);
        gltf_codec.faces = Some(faces_na);
        gltf_codec.uvs = Some(uvs_na);
        let smpl_joints = smpl_output.joints.clone().to_ndarray();
        let joint_count = smpl_joints.shape()[0];
        let lbs_weights = smpl_model.lbs_weights_split().to_ndarray();
        let vertex_count = smpl_output.verts.dims()[0];
        let mut skin_vertex_index = DMatrix::<u32>::zeros(vertex_count, 4);
        let mut skin_vertex_weight = DMatrix::<f32>::zeros(vertex_count, 4);
        for (vertex_id, row) in lbs_weights.outer_iter().enumerate() {
            let mut vertex_weights: Vec<(usize, f32)> = row.iter().enumerate().map(|(index, &weight)| (index, weight)).collect();
            vertex_weights.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
            assert_eq!(vertex_weights.len().min(4), 4, "Illegal vertex weights");
            for (i, (index, weight)) in vertex_weights.iter().take(4).enumerate() {
                skin_vertex_index[(vertex_id, i)] = u32::try_from(*index).expect("Cannot convert to u32!");
                skin_vertex_weight[(vertex_id, i)] = *weight;
            }
        }
        gltf_codec.joint_index = Some(skin_vertex_index);
        gltf_codec.joint_weight = Some(skin_vertex_weight);
        let diffuse_img = scene.get_comp::<&DiffuseImg>(&entity);
        if let Ok(diffuse_img) = diffuse_img {
            if let Some(img) = &diffuse_img.generic_img.cpu_img {
                current_body.diffuse_textures = Some(get_image(img, false, options.max_texture_size));
            }
        }
        let normals_img = scene.get_comp::<&NormalImg>(&entity);
        if let Ok(normals_img) = normals_img {
            if let Some(img) = &normals_img.generic_img.cpu_img {
                current_body.normals_textures = Some(get_image(img, false, options.max_texture_size));
            }
        }
        let metalness_img = scene.get_comp::<&MetalnessImg>(&entity);
        if let Ok(metalness_img) = metalness_img {
            if let Some(img) = &metalness_img.generic_img.cpu_img {
                current_body.metalness_textures = Some(get_image(img, true, options.max_texture_size));
            }
        }
        let roughness_img = scene.get_comp::<&RoughnessImg>(&entity);
        if let Ok(roughness_img) = roughness_img {
            if let Some(img) = &roughness_img.generic_img.cpu_img {
                current_body.roughness_textures = Some(get_image(img, true, options.max_texture_size));
            }
        }
        if scene.world.has::<Pose>(entity).unwrap() && !scene.world.has::<Animation>(entity).unwrap() {
            let Ok(pose_ref) = scene.get_comp::<&Pose>(&entity) else {
                panic!("Pose component doesn't exist");
            };
            let current_pose: &Pose = &pose_ref;
            let current_body_translation = current_pose.global_trans.to_ndarray().to_shape((1, 3)).unwrap().to_owned();
            current_body.joint_poses = Some(current_pose.joint_poses.clone().to_ndarray());
            current_body.body_translation = Some(current_body_translation);
            if smpl_params.enable_pose_corrective {
                let vertex_offsets_merged = smpl_model.compute_pose_correctives(current_pose).to_ndarray();
                let mapping = &smpl_model.idx_split_2_merged_vec();
                let cols = vec![0, 1, 2];
                let vertex_offsets = vertex_offsets_merged.gather(mapping, &cols).into_nalgebra();
                current_body.positions = Some(current_body.positions.as_ref().unwrap() + vertex_offsets);
            }
        }
        #[allow(clippy::cast_precision_loss)]
        if scene.world.has::<Animation>(entity).unwrap() {
            let scene_anim = scene.get_resource::<&SceneAnimation>().unwrap();
            nr_frames = scene_anim.num_frames;
            info!("Processing Animation for body {body_idx:?}");
            let anim = scene.get_comp::<&Animation>(&entity).unwrap();
            let mut current_body_rotations = nd::Array3::<f32>::zeros((joint_count, nr_frames, 3));
            let mut current_body_translations = nd::Array2::<f32>::zeros((nr_frames, 3));
            let mut current_body_scales = nd::Array2::<f32>::zeros((nr_frames, 3));
            let mut current_per_frame_blend_weights = nd::Array2::<f32>::zeros((nr_frames, num_total_blendshapes));
            if should_export_posedirs || should_export_exprdirs {
                let mut full_morph_targets = nd::Array3::<f32>::zeros((num_total_blendshapes, vertex_count, 3));
                let mut running_idx_morph_target = 0;
                if should_export_posedirs {
                    let mut pose_morph_targets = nd::Array3::<f32>::zeros((metadata.num_pose_blend_shapes + 1, vertex_count, 3));
                    let nr_elem_merged = smpl_model.get_pose_dirs().dims()[0] / 3;
                    let pose_dirs_merged = smpl_model
                        .get_pose_dirs()
                        .to_ndarray()
                        .into_shape_with_order((nr_elem_merged, 3, metadata.num_pose_blend_shapes))
                        .unwrap();
                    let mapping = smpl_model.idx_split_2_merged_vec();
                    let cols = vec![0, 1, 2];
                    let depth = (0..metadata.num_pose_blend_shapes).collect::<Vec<_>>().into_boxed_slice();
                    let pose_blend_shapes = pose_dirs_merged
                        .gather(mapping, &cols, &depth)
                        .into_shape_with_order((vertex_count, 3, metadata.num_pose_blend_shapes))
                        .unwrap()
                        .permuted_axes([2, 0, 1]);
                    let pi = nd::Array1::<f32>::from_elem(metadata.num_pose_blend_shapes, -PI);
                    let pi_array = pi.insert_axis(nd::Axis(1)).insert_axis(nd::Axis(2));
                    assert_eq!(pose_blend_shapes.shape()[0], pi_array.len());
                    let template_offset = (pose_blend_shapes.clone() * &pi_array).sum_axis(nd::Axis(0));
                    pose_morph_targets
                        .slice_mut(s![0..metadata.num_pose_blend_shapes, .., ..])
                        .assign(&(pose_blend_shapes));
                    pose_morph_targets
                        .slice_mut(s![metadata.num_pose_blend_shapes, .., ..])
                        .assign(&template_offset);
                    #[allow(clippy::range_plus_one)]
                    full_morph_targets
                        .slice_mut(s![
                            running_idx_morph_target..running_idx_morph_target + metadata.num_pose_blend_shapes + 1,
                            ..,
                            ..
                        ])
                        .assign(&pose_morph_targets);
                    running_idx_morph_target += metadata.num_pose_blend_shapes + 1;
                    gltf_codec.num_pose_morph_targets = metadata.num_pose_blend_shapes + 1;
                }
                #[allow(unused_assignments)]
                if should_export_exprdirs {
                    if let Some(expr_dirs) = smpl_model.get_expression_dirs() {
                        let mut expression_morph_targets = nd::Array3::<f32>::zeros((num_expression_blend_shapes + 1, vertex_count, 3));
                        let nr_elem_merged = expr_dirs.dims()[0] / 3;
                        let expression_dirs_merged = expr_dirs
                            .to_ndarray()
                            .into_shape_with_order((nr_elem_merged, 3, num_expression_blend_shapes))
                            .unwrap();
                        let mapping = smpl_model.idx_split_2_merged_vec();
                        let cols = vec![0, 1, 2];
                        let depth = (0..num_expression_blend_shapes).collect::<Vec<_>>().into_boxed_slice();
                        let expression_dirs_split = expression_dirs_merged
                            .gather(mapping, &cols, &depth)
                            .into_shape_with_order((vertex_count, 3, num_expression_blend_shapes))
                            .unwrap()
                            .permuted_axes([2, 0, 1]);
                        let expression_bounds = nd::Array1::<f32>::from_elem(num_expression_blend_shapes, -7.0);
                        let expression_bounds_array = expression_bounds.insert_axis(nd::Axis(1)).insert_axis(nd::Axis(2));
                        assert_eq!(expression_dirs_split.shape()[0], expression_bounds_array.len());
                        let template_offset = (expression_dirs_split.clone() * &expression_bounds_array).sum_axis(nd::Axis(0));
                        expression_morph_targets
                            .slice_mut(s![0..num_expression_blend_shapes, .., ..])
                            .assign(&(expression_dirs_split));
                        expression_morph_targets
                            .slice_mut(s![num_expression_blend_shapes, .., ..])
                            .assign(&template_offset);
                        #[allow(clippy::range_plus_one)]
                        full_morph_targets
                            .slice_mut(s![
                                running_idx_morph_target..running_idx_morph_target + num_expression_blend_shapes + 1,
                                ..,
                                ..
                            ])
                            .assign(&expression_morph_targets);
                        running_idx_morph_target += num_expression_blend_shapes + 1;
                        gltf_codec.num_expression_morph_targets = num_expression_blend_shapes + 1;
                    }
                }
                gltf_codec.morph_targets = Some(full_morph_targets);
            }
            for global_frame_idx in 0..nr_frames {
                if global_frame_idx < anim.start_offset || global_frame_idx > anim.start_offset + anim.num_animation_frames() {
                    continue;
                }
                let mut local_frame_idx = global_frame_idx - anim.start_offset;
                if global_frame_idx == (anim.start_offset + anim.num_animation_frames()) {
                    local_frame_idx -= 1;
                }
                let mut pose = anim.get_pose_at_idx(local_frame_idx);
                let pose_remap = PoseRemap::new(pose.smpl_type, smpl_params.smpl_type);
                pose = pose_remap.remap(&pose);
                if let Ok(ref pose_mask) = scene.get_comp::<&PoseOverride>(&entity) {
                    let mut new_pose_mask = PoseOverride::clone(pose_mask);
                    pose.apply_mask(&mut new_pose_mask);
                }
                if let Ok(ref pose_retarget) = scene.get_comp::<&RetargetPoseYShift>(&entity) {
                    let mut pose_retarget_local = RetargetPoseYShift::clone(pose_retarget);
                    pose_retarget_local.apply(&mut pose);
                }
                current_body_rotations
                    .slice_mut(s![.., global_frame_idx, ..])
                    .assign(&pose.joint_poses.to_ndarray());
                let mut skeleton_root_translation = pose.global_trans.to_ndarray().to_owned();
                let root_translation = smpl_output.joints.to_ndarray().slice(s![0, ..]).to_owned();
                skeleton_root_translation = skeleton_root_translation + root_translation;
                current_body_translations
                    .slice_mut(s![global_frame_idx, ..])
                    .assign(&skeleton_root_translation);
                if global_frame_idx < (anim.start_offset + anim.num_animation_frames()) {
                    current_body_scales.slice_mut(s![global_frame_idx, ..]).assign(&nd::Array1::ones(3));
                }
                let mut running_idx_morph_target = 0;
                if should_export_posedirs {
                    let pose_blend_weights = &smpl_model.compute_pose_feature(&pose);
                    current_per_frame_blend_weights
                        .slice_mut(s![global_frame_idx, 0..metadata.num_pose_blend_shapes])
                        .assign(&pose_blend_weights.to_ndarray());
                    if global_frame_idx == (anim.start_offset + anim.num_animation_frames()) {
                        current_per_frame_blend_weights
                            .slice_mut(s![global_frame_idx..nr_frames, 0..metadata.num_pose_blend_shapes])
                            .assign(&pose_blend_weights.to_ndarray());
                    }
                    running_idx_morph_target += metadata.num_pose_blend_shapes + 1;
                }
                #[allow(unused_assignments)]
                if should_export_exprdirs {
                    let expr_opt = anim.get_expression_at_idx(local_frame_idx);
                    if let Some(expr) = expr_opt.as_ref() {
                        let expr_nd = expr.expr_coeffs.to_ndarray();
                        let max_nr_expr_coeffs = num_expression_blend_shapes.min(expr_nd.len());
                        let expr_coeffs = expr_nd.slice(s![0..max_nr_expr_coeffs]);
                        current_per_frame_blend_weights
                            .slice_mut(s![
                                global_frame_idx,
                                running_idx_morph_target..running_idx_morph_target + max_nr_expr_coeffs
                            ])
                            .assign(&expr_coeffs);
                    }
                    running_idx_morph_target += num_expression_blend_shapes + 1;
                }
            }
            current_body.body_scales = Some(current_body_scales);
            current_body.body_translations = Some(current_body_translations);
            current_body.body_rotations = Some(current_body_rotations);
            if should_export_posedirs {
                current_per_frame_blend_weights
                    .slice_mut(s![.., metadata.num_pose_blend_shapes])
                    .assign(&nd::Array1::<f32>::from_elem(nr_frames, 0.0));
            }
            if should_export_exprdirs {
                current_per_frame_blend_weights
                    .slice_mut(s![.., metadata.num_pose_blend_shapes + num_expression_blend_shapes + 1])
                    .assign(&nd::Array1::<f32>::from_elem(nr_frames, 0.0));
            }
            if should_export_posedirs || should_export_exprdirs {
                current_body.per_frame_blend_weights = Some(current_per_frame_blend_weights);
            }
        }
        current_body.default_joint_translations = Some(smpl_joints);
        gltf_codec.per_body_data.push(current_body);
    }
    let mut query_props = scene
        .world
        .query::<(&Verts, &Faces, &TransformSequence, Option<&UVs>, Option<&DiffuseImg>, &Name)>();
    for (_entity, (verts, faces, transform_sequence, uv, diffuse_img, _name)) in query_props.iter() {
        let scales_vec3 = transform_sequence.scales.iter().flat_map(|s| vec![*s, *s, *s]).collect::<Vec<f32>>();
        let scales_ndarray3 = nd::Array2::from_shape_vec((transform_sequence.num_frames(), 3), scales_vec3).unwrap();
        let prop_data = PropData {
            positions: verts.0.to_ndarray().into_nalgebra(),
            faces: faces.0.to_ndarray().into_nalgebra(),
            normals: None,
            uvs: uv.map(|uv| uv.0.to_ndarray().into_nalgebra()),
            diffuse_texture: diffuse_img.and_then(|diffuse_img| {
                diffuse_img
                    .generic_img
                    .cpu_img
                    .as_ref()
                    .map(|img| get_image(img, false, options.max_texture_size))
            }),
            metalness_texture: None,
            roughness_texture: None,
            normals_texture: None,
            translations: transform_sequence.translations.clone(),
            rotations: transform_sequence.rotations.clone().insert_axis(nd::Axis(0)),
            scales: scales_ndarray3.clone(),
            default_translation: nd::Array2::<f32>::zeros([1, 3]),
            default_joint_poses: nd::Array2::<f32>::zeros([1, 3]),
        };
        gltf_codec.props.push(prop_data);
    }
    info!(
        "Writing {} body scene to GltfCodec: Took {} seconds for {} frames",
        num_bodies,
        now.elapsed().as_secs(),
        nr_frames
    );
    gltf_codec
}
