/********************************************************************/
/***** GPU Wave Front Fetch Graph Cut *******************************/
/********************************************************************/
////////////////////////////////////////////////////
// Copyright (c) 2018 Kiyoshi Oguri    2018.07.07 //
// Released under the MIT license                 //
// http://opensource.org/licenses/mit-license.php //
////////////////////////////////////////////////////
#include <stdio.h>
#include <chrono>

extern int cost(int S, int H, int W);
extern void cut(int S, int H, int W);

#define TD_NUM (256)
int LOOP = 6000;

int SIZE_S;
int SIZE_H;
int SIZE_W;
int SIZE_HW;
int SIZE_SHW;

__constant__ int size_s;
__constant__ int size_h;
__constant__ int size_w;
__constant__ int size_hw;
__constant__ int size_shw;

int *h_FLW;
int *h_OVF;
int *h_TIM;

int *d_FLW;
int *d_OVF;
int *d_TIM;

__constant__ int *FLW;
__constant__ int *OVF;
__constant__ int *TIM;

size_t FLW_SIZE;
size_t OVF_SIZE;
size_t TIM_SIZE;

#define h_ADR1(S, H, W) \
 ((S)*SIZE_HW+\
  (H)*SIZE_W+\
  (W))

#define h_ADR2(S, H, W, D) \
 ((D)*SIZE_SHW+\
  (S)*SIZE_HW+\
  (H)*SIZE_W+\
  (W))

#define ADR1(S, H, W) \
 ((S)*size_hw+\
  (H)*size_w+\
  (W))

#define ADR2(S, H, W, D) \
 ((D)*size_shw+\
  (S)*size_hw+\
  (H)*size_w+\
  (W))

inline __device__ int edg_read(int S, int H, int W, int D) {
  return FLW[ADR2(S, H, W, D)];
}

inline __device__ void edg_add(int S, int H, int W, int D, int V) {
  FLW[ADR2(S, H, W, D)] += V;
}

inline __device__ int ovf_read(int S, int H, int W) {
  return OVF[ADR1(S, H, W)];
}

inline __device__ void ovf_add(int S, int H, int W, int V) {
  atomicAdd(&(OVF[ADR1(S, H, W)]), V);
}

inline __device__ int tim_read(int S, int H, int W) {
  return TIM[ADR1(S, H, W)];
}

inline __device__ void tim_write(int S, int H, int W, int V) {
  TIM[ADR1(S, H, W)] = V;
}

__global__ void wave_init(void) {
  ///////////////////////////////
  int total_id = blockDim.x * blockIdx.x + threadIdx.x;
  if (total_id >= size_shw) return;
  int S  = total_id / size_hw;
  int sa = total_id % size_hw;
  int H  = sa       / size_w;
  int W  = sa       % size_w;
  ///////////////////////////////
  tim_write(S, H, W, 0);
}

__device__ void push(int S, int H, int W, int nS, int nH, int nW, int D, int R) {
  int tim = tim_read(S, H, W);
  int Tim = tim_read(nS, nH, nW);
  if (Tim <= tim) return;
  int edg = edg_read(S, H, W, D);
  if (edg == 0) return;
  int ovf = ovf_read(S, H, W);
  if (ovf > 0) {
    bool q = edg > ovf;
    int pp = (q)? ovf: edg;
    ovf_add(nS, nH, nW, pp);
    ovf_add(S, H, W, -pp);
    edg_add(nS, nH, nW, R, pp);
    edg_add(S, H, W, D, -pp);
    if (q) tim_write(S, H, W, Tim);
  }
  else tim_write(S, H, W, Tim);
}

__global__ void wave(int tt) {
  ///////////////////////////////
  int total_id = blockDim.x * blockIdx.x + threadIdx.x;
  if (total_id >= size_shw) return;
  int S  = total_id / size_hw;
  int sa = total_id % size_hw;
  int H  = sa       / size_w;
  int W  = sa       % size_w;
  ///////////////////////////////
  if (ovf_read(S, H, W) < 0) {
    tim_write(S, H, W, tt);
    return;
  }
  if  (S!=size_s-1)                 push(S, H, W, S+1, H,   W,   0, 9);
  if ((S!=size_s-1)&&(W!=size_w-1)) push(S, H, W, S+1, H,   W+1, 5, 7);
  if ((S!=size_s-1)&&(W!=0))        push(S, H, W, S+1, H,   W-1, 6, 8);
  if                 (W!=size_w-1)  push(S, H, W, S,   H,   W+1, 4, 3);
  if                 (W!=0)         push(S, H, W, S,   H,   W-1, 3, 4);
  if                 (H!=size_h-1)  push(S, H, W, S,   H+1, W,   2, 1);
  if                 (H!=0)         push(S, H, W, S,   H-1, W,   1, 2);
  if ((S!=0)       &&(W!=size_w-1)) push(S, H, W, S-1, H,   W+1, 8, 6);
  if ((S!=0)       &&(W!=0))        push(S, H, W, S-1, H,   W-1, 7, 5);
  if  (S!=0)                        push(S, H, W, S-1, H,   W,   9, 0);
}

void wave_cut(int loop) {
  wave_init<<< (SIZE_SHW/TD_NUM)+1, TD_NUM >>>();
  for (int time = 1; time <= loop; time++) {
    wave<<< (SIZE_SHW/TD_NUM)+1, TD_NUM >>>(time);
  }
}

void data_set(int penalty_w, int penalty_h, int inhibit_a, int inhibit_b) {
  for (int H = 0; H < SIZE_H; H++) {
    for (int W = 0; W < SIZE_W; W++) {
      for (int S = 0; S < SIZE_S; S++) {
        ///////////////////////////////
        for (int i = 0; i < 10; i++) h_FLW[h_ADR2(S, H, W, i)] = 0;
        h_OVF[h_ADR1(S, H, W)] = 0;
        ///////////////////////////////
        if  (S==SIZE_S-1)          h_OVF[h_ADR1(S, H, W)]   -= cost(S+1, H, W);
        if  (S==0)                 h_OVF[h_ADR1(S, H, W)]    = cost(S,   H, W);
        if  (S!=SIZE_S-1)          h_FLW[h_ADR2(S, H, W, 0)] = cost(S+1, H, W);
        if          (W!=SIZE_W-1)  h_FLW[h_ADR2(S, H, W, 4)] = penalty_w;
        if          (W!=0)         h_FLW[h_ADR2(S, H, W, 3)] = penalty_w;
        if          (H!=SIZE_H-1)  h_FLW[h_ADR2(S, H, W, 2)] = penalty_h;
        if          (H!=0)         h_FLW[h_ADR2(S, H, W, 1)] = penalty_h;
        if ((S!=0)&&(W!=SIZE_W-1)) h_FLW[h_ADR2(S, H, W, 8)] = inhibit_b;
        if ((S!=0)&&(W!=0))        h_FLW[h_ADR2(S, H, W, 7)] = inhibit_b;
        if  (S!=0)                 h_FLW[h_ADR2(S, H, W, 9)] = inhibit_a;
      }
    }
  }
  cudaMemcpy(d_FLW, h_FLW, FLW_SIZE, cudaMemcpyHostToDevice);
  cudaMemcpy(d_OVF, h_OVF, OVF_SIZE, cudaMemcpyHostToDevice);
}

int sink_chk(void) {
  cudaMemcpy(h_OVF, d_OVF, OVF_SIZE, cudaMemcpyDeviceToHost);
  int total = 0;
  for (int H = 0; H < SIZE_H; H++) {
    for (int W = 0; W < SIZE_W; W++) {
      int sink = h_OVF[h_ADR1(SIZE_S-1, H, W)];
      if (sink < 0) total += -sink;
    }
  }
  return total;
}

void dep_set(void) {
  cudaMemcpy(h_TIM, d_TIM, TIM_SIZE, cudaMemcpyDeviceToHost);
  for (int H = 0; H < SIZE_H; H++) {
    for (int W = 0; W < SIZE_W; W++) {
      for (int S = SIZE_S; S >= 0; S--) {
        if (S == 0) cut(S, H, W);
        else if (S == SIZE_S) {
          if (LOOP - h_TIM[h_ADR1(S-1, H, W)] > LOOP/10) {
            cut(S, H, W);
            break;
          }
        }
        else {
          if (h_TIM[h_ADR1(S, H, W)] - h_TIM[h_ADR1(S-1, H, W)] > LOOP/10) {
            cut(S, H, W);
            break;
          }
        }
      }
    }
  }
}

int graph_cut(int penalty_w, int penalty_h, int inhibit_a, int inhibit_b) {
  printf("S-H-W=%d-%d-%d LOOP=%d\n", SIZE_S, SIZE_H, SIZE_W, LOOP);
  cudaMemcpyToSymbol(size_s,   &SIZE_S,   sizeof(int));
  cudaMemcpyToSymbol(size_h,   &SIZE_H,   sizeof(int));
  cudaMemcpyToSymbol(size_w,   &SIZE_W,   sizeof(int));
  cudaMemcpyToSymbol(size_hw,  &SIZE_HW,  sizeof(int));
  cudaMemcpyToSymbol(size_shw, &SIZE_SHW, sizeof(int));
  FLW_SIZE = sizeof(int) * SIZE_SHW*10;
  OVF_SIZE = sizeof(int) * SIZE_SHW;
  TIM_SIZE = sizeof(int) * SIZE_SHW;
  h_FLW = new int[SIZE_SHW*10];
  h_OVF = new int[SIZE_SHW];
  h_TIM = new int[SIZE_SHW];
  cudaMalloc((void **)&d_FLW, FLW_SIZE);
  cudaMalloc((void **)&d_OVF, OVF_SIZE);
  cudaMalloc((void **)&d_TIM, TIM_SIZE);
  cudaMemcpyToSymbol(FLW, &d_FLW, sizeof(int*));
  cudaMemcpyToSymbol(OVF, &d_OVF, sizeof(int*));
  cudaMemcpyToSymbol(TIM, &d_TIM, sizeof(int*));
  data_set(penalty_w, penalty_h, inhibit_a, inhibit_b);
  //---------------------------//
  int before = sink_chk();
  auto start = std::chrono::system_clock::now();
  wave_cut(LOOP);
  auto end = std::chrono::system_clock::now();
  auto dur = end - start;
  int usec = std::chrono::duration_cast<std::chrono::microseconds>(dur).count();
  printf("wave time = %dus\n", usec);
  int after = sink_chk();
  printf("wave flow = %d\n", before - after);
  //---------------------------//
  dep_set();
  cudaFree(d_TIM);
  cudaFree(d_OVF);
  cudaFree(d_FLW);
  delete [] h_TIM;
  delete [] h_OVF;
  delete [] h_FLW;
  return before - after;
}