from __future__ import print_function
import time
import random
cimport cython
import numpy as np
cimport numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

from numpy import int32
from numpy cimport int32_t

from libcpp cimport bool
from libc.math cimport sqrt, cos, sin, pi, exp
from libc.stdlib cimport rand, RAND_MAX, srand

srand(420)

cdef float prec = 10 ** 3
cdef int picture = 350
cdef int matrix_size = 1000
ctypedef np.uint8_t uint8

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
def slope(float I, float E, int rivers, float D):
  cdef int x, y, i, j, k, step
  cdef float dif, prob_sum, prob_0, direction, temp

  cdef int x_dim = 300
  cdef int y_dim = 200

  cdef float prob[4]
  cdef float prob_partial[4]

  start = time.time()
  ims = []
  fig = plt.figure()

  cdef int moves[4][2]
  moves[0][:] = [1, 0]
  moves[1][:] = [0, 1]
  moves[2][:] = [-1, 0]
  moves[3][:] = [0, -1]

  grid_h0 = np.zeros((300, 200), dtype = float)
  for i in range(300):
    grid_h0[i][:] = (<float>i) * I
  cdef double[:, :] grid_h = grid_h0
#  print(grid_h0)

  grid_dh_x0 = np.zeros((300, 200), dtype = np.float)
  grid_dh_y0 = np.zeros((300, 200), dtype = np.float)
  cdef double[:, :] grid_dh_x = grid_dh_x0
  cdef double[:, :] grid_dh_y = grid_dh_y0

  grid_m0 = np.zeros((300, 200), dtype = np.float)
  cdef double[:, :] grid_m = grid_m0


  for i in range(rivers):
    for j in range(50):
      #grid_dh_x = grid_h - np.roll(grid_h, + 1, axis = 0) #difference in x dimension between
      #grid_dh_y = grid_h - np.roll(grid_h, + 1, axis = 1)

      x = rand() % 300
      y = rand() % 200
      grid_m[x][y] = 1
      print(grid_m)
      while x > 2:
        print(x, " ", y)
        prob[:] = [0, 0, 0, 0]
        prob_partial[:] = [0, 0, 0, 0]
        prob_sum = 0

        prob[0] = grid_h[x][y] - grid_h[(x + 1) % x_dim][y]
        prob[1] = grid_h[x][y] - grid_h[x][(y + 1) % y_dim]
        prob[2] = grid_h[x][y] - grid_h[(x - 1) % x_dim][y]
        prob[3] = grid_h[x][y] - grid_h[x][(y - 1) % y_dim]
        print(prob)
        for k in range(4):
          temp = 0
          prob_partial[k] = prob_sum
          if prob[k] >= 0:
            temp = exp(E * prob[k])
          prob_sum += temp

        # if grid_dh_x[x][y] >= 0:
        #   prob_sum += exp(E * grid_dh_x[x][y])
        # if grid_dh_y[x][y] >= 0:
        #   prob[1] = prob_sum
        #   prob_sum += exp(E * grid_dh_y[x][y])
        # if grid_dh_x[x - 1][y] <= 0:
        #   prob[2] = prob_sum
        #   prob_sum += exp(-E * grid_dh_x[x - 1][y])
        # if grid_dh_y[x][y - 1] <= 0:
        #   prob[3] = prob_sum
        #   prob_sum += exp(-E * grid_dh_x[x][y - 1])
        # print(prob)
        # print(prob_sum)
        step = 0
        direction = (<float>rand() / RAND_MAX) * prob_sum
        for k in range(4):
          if direction > prob_partial[k]:
            step = k
        x = (x + moves[step][0]) % x_dim
        y = (y + moves[step][1]) % y_dim
        #print(x, y)
        grid_m[x][y] = 1


      grid_h0 -= D * grid_m0
      grid_m[:][:] = 0

    print(grid_h0)
    plt.matshow(grid_h0)
    plt.show()


  #   if particles % 50 == 0:
  #     print(" " + str(format(time.time() - start, ".1f")) + "s: " + str(particles) + "/" + str(max_par), end = '\r')
  #     ttl = plt.text(450, 0, str(particles) + "/" + str(max_par), horizontalalignment = 'center', verticalalignment = 'bottom')
  #     im = plt.imshow(grid0[offset - picture:offset + picture, offset - picture:offset + picture], animated = True, interpolation = "none", extent = [-picture, picture, -picture, picture])
  #     ims.append([im, ttl], )
  #
  # print("Making an animation...")
  # ani = animation.ArtistAnimation(fig, ims, interval = 50, blit = True, repeat_delay = 100000)
  # plt.title("DLA growth for " + str(max_par) + " particles with probability of 1/" + str(probability) + " to stick")
  # ani.save(str(max_par) + "/" + str(max_par) + "_growth_" + str(probability) + ".mp4", dpi = 150)
  #
  # plt.matshow(grid[offset - picture:offset + picture, offset - picture:offset + picture], extent = [-picture, picture, -picture, picture])
  # plt.title("DLA growth for " + str(max_par) + " particles with probability of 1/" + str(probability) + " to stick", fontsize = '10')
  # plt.savefig(str(max_par) + "/" + str(max_par) + "_" + str(probability) + ".png", dpi = 600)
  # plt.close()
  # exec_time = time.time() - start
  # print("Calculated " + str(max_par) + " particle aggregation with probability of 1/" + str(probability) + " to stick in " + str(format(exec_time, ".1f")) + "s")
  # return 0