#include "kernel_pinhole_fixed_pose_res_jac_first.h"
#include "memops.cuh"
#include <cooperative_groups.h>
#include <cooperative_groups/details/partitioning.h>
#include <cooperative_groups/memcpy_async.h>
#include <cooperative_groups/reduce.h>
#include <cuda_runtime.h>

namespace cg = cooperative_groups;

namespace caspar {

__global__ void __launch_bounds__(1024, 1) PinholeFixedPoseResJacFirstKernel(
    double* sensor_from_rig,
    unsigned int sensor_from_rig_num_alloc,
    double* calib,
    unsigned int calib_num_alloc,
    SharedIndex* calib_indices,
    double* point,
    unsigned int point_num_alloc,
    SharedIndex* point_indices,
    double* pixel,
    unsigned int pixel_num_alloc,
    double* pose,
    unsigned int pose_num_alloc,
    double* out_res,
    unsigned int out_res_num_alloc,
    double* const out_rTr,
    double* out_calib_jac,
    unsigned int out_calib_jac_num_alloc,
    double* const out_calib_njtr,
    unsigned int out_calib_njtr_num_alloc,
    double* const out_calib_precond_diag,
    unsigned int out_calib_precond_diag_num_alloc,
    double* const out_calib_precond_tril,
    unsigned int out_calib_precond_tril_num_alloc,
    double* out_point_jac,
    unsigned int out_point_jac_num_alloc,
    double* const out_point_njtr,
    unsigned int out_point_njtr_num_alloc,
    double* const out_point_precond_diag,
    unsigned int out_point_precond_diag_num_alloc,
    double* const out_point_precond_tril,
    unsigned int out_point_precond_tril_num_alloc,
    size_t problem_size) {
  const int global_thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
  __shared__ uint8_t inout_shared[16384];

  __shared__ SharedIndex calib_indices_loc[1024];
  calib_indices_loc[threadIdx.x] =
      (global_thread_idx < problem_size
           ? calib_indices[global_thread_idx]
           : SharedIndex{0xffffffff, 0xffff, 0xffff});
  __shared__ SharedIndex point_indices_loc[1024];
  point_indices_loc[threadIdx.x] =
      (global_thread_idx < problem_size
           ? point_indices[global_thread_idx]
           : SharedIndex{0xffffffff, 0xffff, 0xffff});

  __shared__ double out_rTr_local[1];

  double r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, r12, r13, r14, r15,
      r16, r17, r18, r19, r20, r21, r22, r23, r24, r25, r26, r27, r28, r29, r30,
      r31, r32, r33, r34, r35, r36, r37, r38, r39, r40, r41, r42, r43, r44, r45,
      r46, r47;
  LoadShared<2, double, double>(
      calib, 2 * calib_num_alloc, calib_indices_loc, (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    ReadShared2<double>(
        (double*)inout_shared, calib_indices_loc[threadIdx.x].target, r0, r1);
  };
  __syncthreads();
  if (global_thread_idx < problem_size) {
    ReadIdx2<1024, double, double, double2>(
        pixel, 0 * pixel_num_alloc, global_thread_idx, r2, r3);
    r4 = -1.00000000000000000e+00;
    r2 = fma(r2, r4, r0);
    r0 = 1.00000000000000008e-15;
    ReadIdx1<1024, double, double, double>(
        sensor_from_rig, 6 * sensor_from_rig_num_alloc, global_thread_idx, r5);
  };
  LoadShared<2, double, double>(
      point, 0 * point_num_alloc, point_indices_loc, (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    ReadShared2<double>(
        (double*)inout_shared, point_indices_loc[threadIdx.x].target, r6, r7);
  };
  __syncthreads();
  if (global_thread_idx < problem_size) {
    ReadIdx2<1024, double, double, double2>(sensor_from_rig,
                                            2 * sensor_from_rig_num_alloc,
                                            global_thread_idx,
                                            r8,
                                            r9);
    ReadIdx2<1024, double, double, double2>(
        pose, 2 * pose_num_alloc, global_thread_idx, r10, r11);
    ReadIdx2<1024, double, double, double2>(sensor_from_rig,
                                            0 * sensor_from_rig_num_alloc,
                                            global_thread_idx,
                                            r12,
                                            r13);
    ReadIdx2<1024, double, double, double2>(
        pose, 0 * pose_num_alloc, global_thread_idx, r14, r15);
    r16 = fma(r12, r15, r9 * r10);
    r17 = r13 * r14;
    r16 = fma(r4, r17, r16);
    r16 = fma(r8, r11, r16);
    r17 = 2.00000000000000000e+00;
    r18 = fma(r12, r11, r9 * r14);
    r19 = r8 * r15;
    r18 = fma(r4, r19, r18);
    r18 = fma(r13, r10, r18);
    r19 = r17 * r18;
    r20 = r16 * r19;
    r21 = r12 * r10;
    r21 = fma(r4, r21, r9 * r15);
    r21 = fma(r13, r11, r21);
    r21 = fma(r8, r14, r21);
    r22 = -2.00000000000000000e+00;
    r23 = fma(r13, r15, r12 * r14);
    r23 = fma(r8, r10, r23);
    r23 = fma(r4, r23, r9 * r11);
    r11 = r22 * r23;
    r24 = fma(r21, r11, r20);
    r5 = fma(r6, r24, r5);
    ReadIdx2<1024, double, double, double2>(
        pose, 4 * pose_num_alloc, global_thread_idx, r25, r26);
    r27 = r12 * r8;
    r27 = r27 * r17;
    r28 = r13 * r9;
    r29 = fma(r22, r28, r27);
    ReadIdx1<1024, double, double, double>(
        pose, 6 * pose_num_alloc, global_thread_idx, r30);
    r31 = r12 * r12;
    r31 = r22 * r31;
    r32 = 1.00000000000000000e+00;
    r33 = r13 * r13;
    r33 = fma(r22, r33, r32);
    r34 = r31 + r33;
    r35 = r13 * r8;
    r35 = r35 * r17;
    r36 = r12 * r9;
    r36 = fma(r17, r36, r35);
    r37 = r17 * r16;
    r37 = r37 * r21;
    r38 = fma(r23, r19, r37);
  };
  LoadShared<1, double, double>(
      point, 2 * point_num_alloc, point_indices_loc, (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    ReadShared1<double>(
        (double*)inout_shared, point_indices_loc[threadIdx.x].target, r39);
  };
  __syncthreads();
  if (global_thread_idx < problem_size) {
    r40 = r21 * r21;
    r40 = r22 * r40;
    r41 = r32 + r40;
    r42 = r18 * r18;
    r42 = r42 * r22;
    r41 = r41 + r42;
    r5 = fma(r25, r29, r5);
    r5 = fma(r30, r34, r5);
    r5 = fma(r26, r36, r5);
    r5 = fma(r7, r38, r5);
    r5 = fma(r39, r41, r5);
    r36 = copysign(1.0, r5);
    r36 = fma(r0, r36, r5);
    r0 = 1.0 / r36;
  };
  LoadShared<2, double, double>(
      calib, 0 * calib_num_alloc, calib_indices_loc, (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    ReadShared2<double>(
        (double*)inout_shared, calib_indices_loc[threadIdx.x].target, r5, r34);
  };
  __syncthreads();
  if (global_thread_idx < problem_size) {
    ReadIdx2<1024, double, double, double2>(sensor_from_rig,
                                            4 * sensor_from_rig_num_alloc,
                                            global_thread_idx,
                                            r29,
                                            r43);
    r40 = r32 + r40;
    r44 = r16 * r16;
    r44 = r22 * r44;
    r40 = r40 + r44;
    r29 = fma(r6, r40, r29);
    r19 = r21 * r19;
    r45 = fma(r16, r11, r19);
    r46 = r17 * r21;
    r46 = fma(r23, r46, r20);
    r28 = fma(r17, r28, r27);
    r27 = r8 * r9;
    r20 = r12 * r13;
    r20 = r20 * r17;
    r27 = fma(r22, r27, r20);
    r47 = r8 * r8;
    r47 = r22 * r47;
    r33 = r47 + r33;
    r29 = fma(r7, r45, r29);
    r29 = fma(r39, r46, r29);
    r29 = fma(r30, r28, r29);
    r29 = fma(r26, r27, r29);
    r29 = fma(r25, r33, r29);
    r33 = r5 * r29;
    r2 = fma(r0, r33, r2);
    r3 = fma(r3, r4, r1);
    r1 = r17 * r16;
    r1 = fma(r23, r1, r19);
    r6 = fma(r6, r1, r43);
    r43 = r8 * r9;
    r43 = fma(r17, r43, r20);
    r47 = r32 + r47;
    r47 = r47 + r31;
    r31 = r12 * r9;
    r31 = fma(r22, r31, r35);
    r11 = fma(r18, r11, r37);
    r44 = r32 + r44;
    r44 = r44 + r42;
    r6 = fma(r25, r43, r6);
    r6 = fma(r26, r47, r6);
    r6 = fma(r30, r31, r6);
    r6 = fma(r39, r11, r6);
    r6 = fma(r7, r44, r6);
    r7 = r34 * r6;
    r3 = fma(r0, r7, r3);
    WriteIdx2<1024, double, double, double2>(
        out_res, 0 * out_res_num_alloc, global_thread_idx, r2, r3);
    r39 = fma(r3, r3, r2 * r2);
  };
  SumStore<double>(out_rTr_local,
                   (double*)inout_shared,
                   0,
                   global_thread_idx < problem_size,
                   r39);
  if (global_thread_idx < problem_size) {
    r39 = r29 * r0;
    r31 = r6 * r0;
    WriteIdx2<1024, double, double, double2>(out_calib_jac,
                                             0 * out_calib_jac_num_alloc,
                                             global_thread_idx,
                                             r39,
                                             r31);
    r30 = r4 * r29;
    r30 = r30 * r2;
    r30 = r30 * r0;
    r47 = r4 * r6;
    r47 = r47 * r3;
    r47 = r47 * r0;
    WriteSum2<double, double>((double*)inout_shared, r30, r47);
  };
  FlushSumShared<2, double>(out_calib_njtr,
                            0 * out_calib_njtr_num_alloc,
                            calib_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    r47 = r4 * r2;
    r30 = r4 * r3;
    WriteSum2<double, double>((double*)inout_shared, r47, r30);
  };
  FlushSumShared<2, double>(out_calib_njtr,
                            2 * out_calib_njtr_num_alloc,
                            calib_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    r29 = r29 * r29;
    r36 = r36 * r36;
    r36 = 1.0 / r36;
    r29 = r29 * r36;
    r6 = r6 * r6;
    r6 = r6 * r36;
    WriteSum2<double, double>((double*)inout_shared, r29, r6);
  };
  FlushSumShared<2, double>(out_calib_precond_diag,
                            0 * out_calib_precond_diag_num_alloc,
                            calib_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    WriteSum2<double, double>((double*)inout_shared, r32, r32);
  };
  FlushSumShared<2, double>(out_calib_precond_diag,
                            2 * out_calib_precond_diag_num_alloc,
                            calib_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    r32 = 0.00000000000000000e+00;
    WriteSum2<double, double>((double*)inout_shared, r32, r39);
  };
  FlushSumShared<2, double>(out_calib_precond_tril,
                            0 * out_calib_precond_tril_num_alloc,
                            calib_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    WriteSum2<double, double>((double*)inout_shared, r31, r32);
  };
  FlushSumShared<2, double>(out_calib_precond_tril,
                            4 * out_calib_precond_tril_num_alloc,
                            calib_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    r32 = r5 * r40;
    r36 = r4 * r36;
    r33 = r36 * r33;
    r32 = fma(r24, r33, r0 * r32);
    r31 = r24 * r36;
    r39 = r34 * r1;
    r39 = fma(r0, r39, r7 * r31);
    WriteIdx2<1024, double, double, double2>(out_point_jac,
                                             0 * out_point_jac_num_alloc,
                                             global_thread_idx,
                                             r32,
                                             r39);
    r31 = r5 * r45;
    r31 = fma(r0, r31, r38 * r33);
    r6 = r34 * r44;
    r29 = r38 * r36;
    r29 = fma(r7, r29, r0 * r6);
    WriteIdx2<1024, double, double, double2>(out_point_jac,
                                             2 * out_point_jac_num_alloc,
                                             global_thread_idx,
                                             r31,
                                             r29);
    r6 = r5 * r46;
    r6 = fma(r0, r6, r41 * r33);
    r33 = r34 * r11;
    r30 = r41 * r36;
    r30 = fma(r7, r30, r0 * r33);
    WriteIdx2<1024, double, double, double2>(
        out_point_jac, 4 * out_point_jac_num_alloc, global_thread_idx, r6, r30);
    r33 = r4 * r2;
    r7 = r4 * r3;
    r7 = fma(r39, r7, r32 * r33);
    r33 = r4 * r2;
    r0 = r4 * r3;
    r0 = fma(r29, r0, r31 * r33);
    WriteSum2<double, double>((double*)inout_shared, r7, r0);
  };
  FlushSumShared<2, double>(out_point_njtr,
                            0 * out_point_njtr_num_alloc,
                            point_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    r0 = r4 * r3;
    r7 = r4 * r2;
    r7 = fma(r6, r7, r30 * r0);
    WriteSum1<double, double>((double*)inout_shared, r7);
  };
  FlushSumShared<1, double>(out_point_njtr,
                            2 * out_point_njtr_num_alloc,
                            point_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    r7 = fma(r39, r39, r32 * r32);
    r0 = fma(r29, r29, r31 * r31);
    WriteSum2<double, double>((double*)inout_shared, r7, r0);
  };
  FlushSumShared<2, double>(out_point_precond_diag,
                            0 * out_point_precond_diag_num_alloc,
                            point_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    r0 = fma(r30, r30, r6 * r6);
    WriteSum1<double, double>((double*)inout_shared, r0);
  };
  FlushSumShared<1, double>(out_point_precond_diag,
                            2 * out_point_precond_diag_num_alloc,
                            point_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    r0 = fma(r32, r31, r39 * r29);
    r39 = fma(r39, r30, r32 * r6);
    WriteSum2<double, double>((double*)inout_shared, r0, r39);
  };
  FlushSumShared<2, double>(out_point_precond_tril,
                            0 * out_point_precond_tril_num_alloc,
                            point_indices_loc,
                            (double*)inout_shared);
  if (global_thread_idx < problem_size) {
    r30 = fma(r29, r30, r31 * r6);
    WriteSum1<double, double>((double*)inout_shared, r30);
  };
  FlushSumShared<1, double>(out_point_precond_tril,
                            2 * out_point_precond_tril_num_alloc,
                            point_indices_loc,
                            (double*)inout_shared);
  SumFlushFinal<double>(out_rTr_local, out_rTr, 1);
}

void PinholeFixedPoseResJacFirst(double* sensor_from_rig,
                                 unsigned int sensor_from_rig_num_alloc,
                                 double* calib,
                                 unsigned int calib_num_alloc,
                                 SharedIndex* calib_indices,
                                 double* point,
                                 unsigned int point_num_alloc,
                                 SharedIndex* point_indices,
                                 double* pixel,
                                 unsigned int pixel_num_alloc,
                                 double* pose,
                                 unsigned int pose_num_alloc,
                                 double* out_res,
                                 unsigned int out_res_num_alloc,
                                 double* const out_rTr,
                                 double* out_calib_jac,
                                 unsigned int out_calib_jac_num_alloc,
                                 double* const out_calib_njtr,
                                 unsigned int out_calib_njtr_num_alloc,
                                 double* const out_calib_precond_diag,
                                 unsigned int out_calib_precond_diag_num_alloc,
                                 double* const out_calib_precond_tril,
                                 unsigned int out_calib_precond_tril_num_alloc,
                                 double* out_point_jac,
                                 unsigned int out_point_jac_num_alloc,
                                 double* const out_point_njtr,
                                 unsigned int out_point_njtr_num_alloc,
                                 double* const out_point_precond_diag,
                                 unsigned int out_point_precond_diag_num_alloc,
                                 double* const out_point_precond_tril,
                                 unsigned int out_point_precond_tril_num_alloc,
                                 size_t problem_size) {
  if (problem_size == 0) {
    return;
  }

  const int n_blocks = (problem_size + 1024 - 1) / 1024;
  PinholeFixedPoseResJacFirstKernel<<<n_blocks, 1024>>>(
      sensor_from_rig,
      sensor_from_rig_num_alloc,
      calib,
      calib_num_alloc,
      calib_indices,
      point,
      point_num_alloc,
      point_indices,
      pixel,
      pixel_num_alloc,
      pose,
      pose_num_alloc,
      out_res,
      out_res_num_alloc,
      out_rTr,
      out_calib_jac,
      out_calib_jac_num_alloc,
      out_calib_njtr,
      out_calib_njtr_num_alloc,
      out_calib_precond_diag,
      out_calib_precond_diag_num_alloc,
      out_calib_precond_tril,
      out_calib_precond_tril_num_alloc,
      out_point_jac,
      out_point_jac_num_alloc,
      out_point_njtr,
      out_point_njtr_num_alloc,
      out_point_precond_diag,
      out_point_precond_diag_num_alloc,
      out_point_precond_tril,
      out_point_precond_tril_num_alloc,
      problem_size);
}

}  // namespace caspar