#include <stdlib.h>

// ---
// Header includes
#include "base/base_inc.h"
#include "os/os_inc.h"


// CUDA headers
#include "kernels.h"

// ---
// .C includes
#include "base/base_inc.c"
#include "os/os_inc.c"
#include "os/os_entry_point.c"


#define grid_file_path_bin "D:\\dev\\eigsol_gpu\\out\\grid.bin"
#define grid_file_path "D:\\dev\\eigsol_gpu\\out\\grid.dat"
#define knotpoints_file_path "D:\\dev\\eigsol_gpu\\out\\knotpoints.dat"
#define bspline_array_file_path "D:\\dev\\eigsol_gpu\\out\\bsplines.dat"

global U32 enable_logging = 1;

// Complex number with double precision
typedef struct Z64 Z64;
struct Z64
{
  F64 re;
  F64 im;
};

typedef struct Grid Grid;
struct Grid 
{
  F64 start;
  F64 end;
  U32 num_steps;
  F64 *points;
};


typedef struct BSplineCtx BSplineCtx;
struct BSplineCtx 
{
  Arena* arena;
  U32 order;
  F64 *knotpoints;
  U32 num_knotpoints;
  U32 num_bsplines;
  U32 num_phys_points;
  F64 *bsplines;
};



//~ Globals
global Grid g_grid = {0};
global BSplineCtx g_bspline_ctx = {0};

//~ Functions

function inline void
logger_debug(String8 msg)
{
  OutputDebugString(msg.str);
}

#define LOG(msg) if(enable_logging) { logger_debug(msg) };


function void
write_array_binary_F64(String8 path_to_file, F64 *values, U32 array_size) 
{
  OS_Handle file_handle = OS_file_open(OS_AccessFlag_Write | OS_AccessFlag_CreateNew, 
      path_to_file);
  { 
    ArenaTemp scratch = scratch_get(0, 0);
    String8List list = {0};
    String8 temp = {0};
    temp.str = (U8*)values;
    temp.size = sizeof(F64)*array_size;
    str8_list_push(scratch.arena, &list, temp);
    OS_file_write(scratch.arena, file_handle, 0, list, 0);
    
    String8List log_list = {0};
    str8_list_push(scratch.arena, &log_list, str8_lit("Wrote binary array data to"));
    str8_list_push(scratch.arena, &log_list, path_to_file);
    StringJoin join = {0};
    join.sep = str8_lit(" ");
    join.post = str8_lit("\n");
    String8 log_msg = str8_list_join(scratch.arena, log_list, &join);
    OutputDebugString(log_msg.str);
  }
  OS_file_close(file_handle);
}

function void
write_string_list_to_file(Arena *arena, String8 path, String8List *list)
{

  OS_Handle file_handle = OS_file_open(OS_AccessFlag_Write | OS_AccessFlag_CreateNew, 
      path);
  OS_file_write(arena, file_handle, 0, *list, 0);

  U32 debug = 1;
  if(debug)
  {
    String8List log_list = {0};
    str8_list_push(arena, &log_list, str8_lit("Wrote array to"));
    str8_list_push(arena, &log_list, path);
    StringJoin join = {0};
    join.sep = str8_lit(" ");
    join.post = str8_lit("\n");
    String8 log_msg = str8_list_join(arena, log_list, &join);
    OutputDebugString(log_msg.str);
  }
  OS_file_close(file_handle);
}


  function void
write_array_F64(String8 path_to_file, F64 *values, U32 array_size, char* fmt) 
{
    ArenaTemp scratch = scratch_get(0, 0);
    String8List list = {0};
    for(U32 i = 0; i < array_size; i++) 
    {
      str8_list_pushf(scratch.arena, &list, fmt, values[i]);
    }
    write_string_list_to_file(scratch.arena, path_to_file, &list);
}

  function F64 
bspline_recursion(F64 x, U32 k, U32 i)
{
  F64 *t = g_bspline_ctx.knotpoints;

  if(k == 1)
  {
    if(i == g_bspline_ctx.num_bsplines-1 && x == g_grid.end) 
    {
      // TODO(anton):
      // This is like a hack to get the last bspline to be 1 at the last point.
      // I dont get how the Cox-de Boor recursion formula can force the last bspline 
      // to unity at the last point, actually. I have to check this.
      return 1.0;
    }
    else if(i < g_bspline_ctx.num_bsplines && (x >= t[i] && x < t[i+1]))
    {
      return 1.0;
    } else {
      return 0.0;
    }
  } 
  else 
  {
    F64 recursion1 = bspline_recursion(x, k-1, i);
    F64 term1_enum = (x - t[i]);
    F64 term1_denom = (t[i+k-1] - t[i]);
    F64 term1 = recursion1 > 0.0 ? (term1_enum/term1_denom)*recursion1 : 0.0;

    F64 recursion2 = bspline_recursion(x, k-1, i+1);
    F64 term2_enum = (t[i+k] - x);
    F64 term2_denom = (t[i+k] - t[i+1]);
    F64 term2 = recursion2 > 0.0 ? (term2_enum/term2_denom)*recursion2 : 0.0;

    return term1 + term2;
  }


}


  function F64 
get_bspline_F64(F64 x_coord, U32 index)
{
  U32 k = g_bspline_ctx.order;
  F64 out = bspline_recursion(x_coord, k, index);
  return out;
}

  function void
set_up_grid(Arena *arena)
{
  g_grid.start = 0.0;
  g_grid.end = 10.0;
  g_grid.num_steps = 100;

  g_grid.points = PushArray(arena, F64, g_grid.num_steps);
  F64 step_size = (g_grid.end-g_grid.start)/(F64)g_grid.num_steps;
  g_grid.points[0] = g_grid.start;
  g_grid.points[g_grid.num_steps-1] = g_grid.end;
  for(U32 i = 1; i < g_grid.num_steps-1; i++)
  {
    g_grid.points[i] = g_grid.points[i-1] + step_size;
  }

}


  function void
set_up_bspline_context(Arena* arena)
{
  // Create knotpoint sequence.
  U32 k = 4;
  U32 N = 14;
  g_bspline_ctx.order = k; 
  g_bspline_ctx.num_knotpoints = N; 
  g_bspline_ctx.num_bsplines = N-k;
  g_bspline_ctx.num_phys_points = N-(2*k)+2; // Remove k points at each end, and then add back the first and last points.
  g_bspline_ctx.arena = arena;
  g_bspline_ctx.knotpoints = PushArray(arena, F64, g_bspline_ctx.num_knotpoints);
  // Set up physical points;
  F64 delta = (g_grid.end-g_grid.start)/(g_bspline_ctx.num_phys_points-1);
  // Set ghost points including first physical
  U32 phys_point_last_index = g_bspline_ctx.num_phys_points + k-1;
  for(U32 i = 0; i < k; i++) 
  {
    g_bspline_ctx.knotpoints[i] = g_grid.start;
  }
  for(U32 i = k; i < phys_point_last_index; i++) 
  {
    g_bspline_ctx.knotpoints[i] = g_bspline_ctx.knotpoints[i-1] + delta;
  }
  // Set the last points
  F64 last_physical = g_grid.end;
  for(U32 i = phys_point_last_index; i < g_bspline_ctx.num_knotpoints; i++)
  {
    g_bspline_ctx.knotpoints[i] = last_physical;
  }
}


  function void
write_bsplines_to_matrix_F64(Arena *arena)
{

  U64 num_bsplines = g_bspline_ctx.num_bsplines;
  U64 k = g_bspline_ctx.order;
  F64 *t = g_bspline_ctx.knotpoints;
  U32 grid_size = g_grid.num_steps;

  // For sanity check we make the first 4 bsplines by hand.
  {
    F64 *bspl0 = PushArray(arena, F64, grid_size);
    F64 *bspl1 = PushArray(arena, F64, grid_size);
    F64 *bspl2 = PushArray(arena, F64, grid_size);
    F64 *bspl3 = PushArray(arena, F64, grid_size);
    F64 *bspl9 = PushArray(arena, F64, grid_size);
    for(U32 i = 0; i < grid_size; i++)
    {
      F64 x = g_grid.points[i];
      bspl0[i] = get_bspline_F64(x, 0);
      bspl1[i] = get_bspline_F64(x, 1);
      bspl2[i] = get_bspline_F64(x, 2);
      bspl3[i] = get_bspline_F64(x, 3);
      bspl9[i] = get_bspline_F64(x, 9);
    }

    F64 test = get_bspline_F64(g_grid.points[grid_size-1], 9);

    write_array_F64(str8_lit("D:\\dev\\eigsol_gpu\\out\\bspline0.dat"), bspl0, grid_size, "%13.6e\n");
    write_array_F64(str8_lit("D:\\dev\\eigsol_gpu\\out\\bspline1.dat"), bspl1, grid_size, "%13.6e\n");
    write_array_F64(str8_lit("D:\\dev\\eigsol_gpu\\out\\bspline2.dat"), bspl2, grid_size, "%13.6e\n");
    write_array_F64(str8_lit("D:\\dev\\eigsol_gpu\\out\\bspline3.dat"), bspl3, grid_size, "%13.6e\n");
    write_array_F64(str8_lit("D:\\dev\\eigsol_gpu\\out\\bspline9.dat"), bspl9, grid_size, "%13.6e\n");
  }

  g_bspline_ctx.bsplines = PushArray(arena, F64, grid_size*num_bsplines);
  for(U32 i = 0; i < g_grid.num_steps; i++)
  {
    for(U32 j = 0; j < num_bsplines; j++)
    {
      U32 index = g_grid.num_steps * i + j;
      g_bspline_ctx.bsplines[index] = get_bspline_F64(g_grid.points[i], j);
    }
  }

  ArenaTemp scratch = scratch_get(0, 0);
  String8 bspline_filename = str8_lit("D:\\dev\\eigsol_gpu\\out\\Bsplines.dat");

  // First line is just the bspline indices.
  String8List first_line_list = {0};
  StringJoin join = {0};
  join.sep = str8_lit("\t\t");
  for(U32 i = 0; i < num_bsplines; i++) 
  {
    str8_list_pushf(scratch.arena, &first_line_list, "%i", i);
  }
  String8 first_line = str8_list_join(scratch.arena, first_line_list, &join);

  String8List bspline_array_list = {0};
  for(U32 i = 0; i < g_grid.num_steps; i++)
  {
    String8List row = {0};
    for(U32 j = 0; j < num_bsplines; j++)
    {
      U32 index = g_grid.num_steps * i + j;
      F64 bspline_value = g_bspline_ctx.bsplines[index];
      str8_list_pushf(scratch.arena, &row, "%13.6e", bspline_value);
    }
    StringJoin bspl_join = {0};
    bspl_join.sep = str8_lit("   ");
    bspl_join.post = str8_lit("\n");
    String8 row_joined = str8_list_join(scratch.arena, row, &bspl_join);
    str8_list_push(scratch.arena, &bspline_array_list, row_joined);
  }


  write_string_list_to_file(scratch.arena, str8_lit(bspline_array_file_path), &bspline_array_list);





}


  function
void EntryPoint() 
{
  OS_InitReceipt os_receipt = OS_init();
  OS_InitGfxReceipt os_gfx_receipt = OS_gfx_init(os_receipt);

  Arena *arena = m_make_arena();

  set_up_grid(arena);

  write_array_binary_F64(str8_lit(grid_file_path_bin), g_grid.points, g_grid.num_steps);
  write_array_F64(str8_lit(grid_file_path), g_grid.points, g_grid.num_steps, "%13.6e\n");

  set_up_bspline_context(arena);


  write_array_F64(str8_lit(knotpoints_file_path), g_bspline_ctx.knotpoints, g_bspline_ctx.num_knotpoints, "%13.6e\n");

  write_bsplines_to_matrix_F64(arena);


}



