# Take a toroidal exodus mesh generated by rotating a reference poloidal plane about an axis
# and decompose it into coordinates on the poloidal plane and a coordinate in the toroidal direction
# Generate a numpy tensor with dimensions (toroidal ID, poloidal ID, DOF, time)
# Assumes that this is run with mpi with the same number of procs as files in the exodus decomposition
#
# PYTHONPATH must include the pygenten subdirectory of a Genten build and a path including exodus3.py
#
# Run command:
#  mpiexec -np <num_procs> python3 torus_to_tensor.py <exodus_base_name> <axis_of_rotation> <num_procs_per_poloidal_plane> <tol>
#  
#  <exodus_base_name>              name of a decomposed exodus file up to and including the ".exo"
#                                  e.g. "simple_torus.exo"
#                                  assumes the existence of files <exodus_base_name>.<num_procs>.n, n = 0, ... , <num_procs>-1 
#  <axis_of_rotation>              axis around which a reference poloidal plane has been rotated to generate the mesh: x, y, or z
#  <num_procs_per_poloidal_plane>  when the tensor is fully constructed how many procs the poloidal plane dimension is decomposed over
#                                  this implies a value for <num_procs_in_toroidal_direction> = <num_procs>/<num_procs_per_poloidal_plane>
#                                  and as such must evenly divide <num_procs>
#                                  furthermore, we require <num_procs_per_poloidal_plane> <= the number of nodes in the reference poloidal plane 
#                                  and <num_procs_in_toroidal_direction> <= the number of poloidal planes in the mesh
#                                  (defaults to <num_procs> if not supplied)
#  <tol>                           optional parameter defining a tolerance for scalar comparisons (defaults to 1.0e-10) 

import os
import sys
import exodus3 as ex
import numpy as np
import pygenten as gt
import pyttb
import pygenten._phys_utils as pu

def tensor_to_exodus(tensor, tensor_node_gids, mesh_base_filename, output_base_filename):

    # copy the mesh into output files
    num_procs = gt.num_procs()
    rank = gt.proc_rank()
    mesh_filename = mesh_base_filename+"."+str(num_procs)+"."+str(rank)
    mesh_file = ex.exodus(mesh_filename, mode="r", array_type='numpy')
    output_filename = output_base_filename+"."+str(num_procs)+"."+str(rank)
    if os.path.exists(output_filename):
      os.remove(output_filename)
    output_file = ex.copy_mesh(mesh_filename,
                               output_filename,
                               exoFromObj=mesh_file,
                               additionalElementAttributes=[],
                               array_type='numpy')      
    var_names = mesh_file.get_node_variable_names()
    num_vars = len(var_names)
    times = mesh_file.get_times()
    num_times =len(times)
    for t in range(len(times)):
      output_file.put_time(t+1, times[t])
    output_file.set_node_variable_number(num_vars)
    for v in range(num_vars):
      output_file.put_node_variable_name(var_names[v],v+1)
    mesh_file.close()

    # store the tensor data in something like a multivector
    tensor_data = tensor.data
    num_theta = tensor_data.shape[0]
    num_poloidal = tensor_data.shape[1]
    num_src_nodes = num_poloidal*num_theta
    src_node_data = np.zeros((num_src_nodes,num_times*num_vars), dtype=np.double)
    src_node_gids = np.zeros(num_src_nodes, dtype=np.longlong)
    for th in range(num_theta):
      for p in range(num_poloidal):
        src_node_gids[num_poloidal*th+p] = tensor_node_gids[th,p]
        for t in range(num_times):
          for v in range(num_vars):
            src_node_data[num_poloidal*th+p,t*num_vars+v] = tensor_data[th,p,v,t]
    
    # map the distributed tensor data to the appropriate rank for the exo file
    dest_node_gids = output_file.get_node_id_map()
    num_dest_nodes = len(dest_node_gids)
    dest_node_data = pu.redistribute_data_across_procs(src_node_gids,src_node_data,dest_node_gids)
    
    # write the tensor data into the exo file
    for t in range(num_times):
      for v in range(num_vars):
        values = np.zeros(num_dest_nodes, dtype=np.double)
        for n in range(num_dest_nodes):
          values[n] = dest_node_data[n,t*num_vars+v]
        output_file.put_node_variable_values(var_names[v],t+1,values)
    output_file.close() 

def torus_to_tensor(base_filename, axis, num_procs_per_poloidal_plane, tol=1.0e-10, rescale_variables=True):
    # error check the processor decomposition
    num_procs = gt.num_procs()
    if num_procs_per_poloidal_plane > 0:
      if num_procs%num_procs_per_poloidal_plane > 0:
        raise Exception('Invalid num_procs_per_poloidal_plane (' + str(num_procs_per_poloidal_plane) + '): must divide num MPI ranks (' + str(num_procs) + ') evenly')

    # get filename for this proc
    rank = gt.proc_rank()
    exo_filename = base_filename+"."+str(num_procs)+"."+str(rank)

    # open the file
    exo_file = ex.exodus(exo_filename, mode="r", array_type='numpy')

    # assign ownership to shared nodes and get lists of the owned lids, gids, and coords
    lids, gids, x, y, z = get_owned_ids_and_coords(exo_file)
    num_nodes = len(lids)

    # transform to r, theta, a coordinates (a is x, y, or z depending on axis of rotation)
    r, theta, a = get_cylindrical_coordinates(x, y, z, axis, tol)

    # get list of all values of theta across all procs
    unique_thetas = collect_unique_values_over_procs(theta, tol)

    # assign theta ids (tids) to each node
    tids = -1*np.ones(num_nodes, dtype=np.longlong) 
    for i in range(num_nodes):
      for j in range(len(unique_thetas)):
        if abs(theta[i]-unique_thetas[j]) < tol:
          tids[i] = j
          break
    if min(tids) < 0:
      raise Exception('Some nodes have not been assigned IDs in the theta direction')

    # the poloidal plane with min theta will be our reference plane
    min_theta = min(unique_thetas)

    # get a list of all nodes on this plane
    ref_lids = []
    for i in range(num_nodes):
      if abs(theta[i] - min_theta) < tol:
         ref_lids.append(i)
    num_ref_nodes = len(ref_lids)

    # check that all nodes have been accounted for and that the processor decomposition works
    total_ref_nodes = pu.global_go_sum(num_ref_nodes)
    total_nodes     = pu.global_go_sum(num_nodes)
    total_thetas    = len(unique_thetas)
    num_theta_procs = -1
    if total_nodes != total_ref_nodes*total_thetas:
      msg = 'Toroidal decomposition failure: total_nodes != total_ref_nodes*total_thetas'
      msg = msg + ' (' + str(total_nodes) + ' != ' + str(total_ref_nodes) + '*' + str(total_thetas) + ')'
      raise Exception(msg)
    if num_procs_per_poloidal_plane:
      if num_procs_per_poloidal_plane > total_ref_nodes:
        msg = 'Invalid num_procs_per_poloidal_plane (' + str(num_procs_per_poloidal_plane) + '): '
        msg = msg + 'must be less than number of nodes per poloidal plane (' + str(total_ref_nodes) +')'
        raise Exception(msg)
      num_theta_procs = num_procs//num_procs_per_poloidal_plane
      if num_theta_procs > total_thetas:
        msg = 'Invalid num_theta_procs (num_procs/num_procs_per_poloidal_plane = ' + str(num_theta_procs) + '): '
        msg = msg + 'must be less than number of poloidal planes in the mesh (' + str(total_thetas) +')'
        raise Exception(msg)

    # get the associated gids and r and a values
    ref_gids = np.zeros(num_ref_nodes, dtype=np.longlong)
    ref_r    = np.zeros(num_ref_nodes, dtype=np.double)
    ref_a    = np.zeros(num_ref_nodes, dtype=np.double)
    for i in range(num_ref_nodes):
      ref_i       = ref_lids[i]
      ref_gids[i] = gids[ref_i]
      ref_r[i]    = r[ref_i]
      ref_a[i]    = a[ref_i]

    # we'll assign unique IDs from 0 to num_ref_nodes-1 to the nodes on the reference plane (rids)
    ref_rids = pu.get_reference_ids(ref_gids)

    # assign rids to each node by matching coordinates to the reference plane
    rids = match_coordinates(r, a, ref_r, ref_a, ref_rids, tol)

    # grab the data for each node
    # store it in something like a multivector
    var_names = exo_file.get_node_variable_names()
    num_vars  = len(var_names)
    num_times = exo_file.num_times()
    node_data = np.zeros((num_nodes,num_times*num_vars), dtype=np.double)
    for t in range(num_times):
      for v in range(num_vars):
        this_data = exo_file.get_node_variable_values(var_names[v],t+1)
        for i in range(num_nodes):
          node_data[i,t*num_vars+v] = this_data[lids[i]]
    exo_file.close()

    # create a composite ID (cid) out of the rids and tids
    # this will be used to redistribute the data across procs
    cids = -1*np.ones(num_nodes, dtype=np.longlong) 
    for i in range(num_nodes):
      cids[i] = tids[i]*total_ref_nodes + rids[i]

    # determine which cids should be on which procs based on the user defined decomposition
    redistributed_cids = -1*np.ones(0, dtype=np.longlong)
    global_blocking = []
    parallel_map = []
    if num_procs_per_poloidal_plane > 0:
      redistributed_cids, global_blocking = distribute_composite_ids_across_procs(num_theta_procs,total_thetas,num_procs_per_poloidal_plane,total_ref_nodes)
    else:
      redistributed_cids = distribute_composite_ids_to_root(num_theta_procs,total_thetas,num_procs_per_poloidal_plane,total_ref_nodes)
    redistributed_num_nodes = len(redistributed_cids)

    # use cids to redistribute data across procs
    redistributed_node_data = pu.redistribute_data_across_procs(cids,node_data,redistributed_cids)
    tensor = np.zeros((0,0,0,0), dtype=np.double)
    if len(redistributed_node_data) == 0:
      return tensor

    # communicate the original exodus GIDs to their redistributed locations
    # we'll use this in mapping tensor data back to the exodus mesh
    redistributed_node_gids = pu.redistribute_ids_across_procs(cids,gids,redistributed_cids)

    # back out tids and rids from redistributed cids
    redistributed_tids = -1*np.ones(redistributed_num_nodes, dtype=np.longlong)
    redistributed_rids = -1*np.ones(redistributed_num_nodes, dtype=np.longlong)
    for i in range(redistributed_num_nodes):
      redistributed_tids[i] = redistributed_cids[i]//total_ref_nodes
      redistributed_rids[i] = redistributed_cids[i]%total_ref_nodes

    # theta and reference ids should now be contiguous on each proc
    start_tid = min(redistributed_tids)
    num_local_tids = max(redistributed_tids) - start_tid + 1
    start_rid = min(redistributed_rids)
    num_local_rids = max(redistributed_rids) - start_rid + 1

    # build the tensor
    tensor = np.zeros((num_local_tids,num_local_rids,num_vars,num_times), dtype=np.double)
    tensor_node_gids = np.zeros((num_local_tids,num_local_rids), dtype=np.longlong)
    for i in range(redistributed_num_nodes):
      tensor_node_gids[redistributed_tids[i]-start_tid, redistributed_rids[i]-start_rid] = redistributed_node_gids[i]
      for v in range(num_vars):
        for t in range(num_times):
          tensor[redistributed_tids[i]-start_tid, redistributed_rids[i]-start_rid, v, t] = redistributed_node_data[i,t*num_vars+v]

    # add var and time dimensions to global blocking
    if num_procs_per_poloidal_plane > 0:
      vt_blocking = np.zeros((2,num_procs+1),dtype=np.longlong)
      vt_blocking[0,1] = num_vars
      vt_blocking[1,1] = num_times
      global_blocking = np.vstack([global_blocking,vt_blocking])
      parallel_map = np.ones((4),dtype=np.longlong)
      parallel_map[0] = num_theta_procs
      parallel_map[1] = num_procs_per_poloidal_plane

    # convert to genten tensor, distribute and return
    X = pyttb.tensor.from_data(tensor)

    # compute mu and stdv for rescaling variables
    mu = np.zeros(num_vars, dtype=np.double)
    stdv = np.ones(num_vars, dtype=np.double)
    mode = 2
    if rescale_variables: 

      print("rescaling")
      # list of modes not including mode
      axis = list(range(X.ndims)) 
      axis.pop(mode)
      axis = tuple(axis)
    
      # compute mean and standard deviation
      mu = np.mean(X.data, axis=axis)
      var = np.var(X.data, axis=axis)
      local_N = X.data.size/mu.size
      N = pu.global_go_sum(int(local_N))
      for i in range(len(mu)):
        mu[i] = pu.global_scalar_sum(mu[i]*local_N)/N
        stdv[i]  = np.sqrt(pu.global_scalar_sum(var[i]*local_N)/N)
      smin = 1.0e-12
      stdv[stdv<smin] = 1

    # compatible shape for mu,stdv for arithmetic with X
    shp = [1 for i in range(X.ndims)]
    shp[mode] = X.shape[mode]
    shp = tuple(shp)
    mt =np.reshape(mu,shp)
    st =np.reshape(stdv,shp)

    Xs = (X.data - mt) / st
    X = pyttb.tensor.from_data(Xs)

    gt_tensor = gt.make_gt_tensor(X)
    gt_dist_tensor, gt_dtc = gt.distribute_tensor(gt_tensor, global_blocking, parallel_map)
    del(gt_tensor)
    return gt_dist_tensor, gt_dtc, tensor_node_gids, mt, st

def get_cylindrical_coordinates(x, y, z, axis, tol):
   """
   Map (x,y,z) coordinates to (r,theta,a) coordinates
   a is the coordinate parallel to the specified axis of rotation ('x', 'y', or 'z')
   """

   num_nodes = len(x)
   r     = np.zeros(num_nodes, dtype=np.double)
   theta = np.zeros(num_nodes, dtype=np.double)
   a     = np.zeros(num_nodes, dtype=np.double)
   b     = np.zeros(num_nodes, dtype=np.double)
   c     = np.zeros(num_nodes, dtype=np.double)

   # a is the coordinate in the direction of the axis of rotation
   # b and c are temps for the other two coordinates
   if axis == 'x':
      a = x
      b = y
      c = z
   elif axis == 'y':
      a = y
      b = x
      c = z
   else:
      a = z
      b = x
      c = y
    
   # do the transformation
   for i in range(num_nodes):
      r[i] = np.sqrt(b[i]*b[i] + c[i]*c[i])
      theta[i] = np.arccos(b[i]/r[i])
      if c[i] < 0.0:
        theta[i] *= -1.0
        theta[i] += 2.0*np.arccos(-1.0)
      if abs(theta[i]-2.0*np.arccos(-1.0)) < tol:
        theta[i] = 0.0
   return r, theta, a

def get_owned_ids_and_coords(exo_file):
   """
   Nodes in an exodus file may overlap between processors. This function assigns unique ownership
   over each node and returns a list of LIDs pointing to the owned nodes as well as the corresponding
   GIDs and coordinates
   """

   # get all owned and shared coordinates and gids
   all_x, all_y, all_z = exo_file.get_coords()
   all_gids = exo_file.get_node_id_map()

   # assign nodes to be uniquely owned by procs
   lids = pu.assign_owned_lids(all_gids)
   num_nodes = len(lids)
   gids = np.zeros(num_nodes, dtype=np.longlong)
   x    = np.zeros(num_nodes, dtype=np.double)
   y    = np.zeros(num_nodes, dtype=np.double)
   z    = np.zeros(num_nodes, dtype=np.double)
   for i in range(num_nodes):
      gids[i] = all_gids[lids[i]]
      x[i]    = all_x[lids[i]]
      y[i]    = all_y[lids[i]]
      z[i]    = all_z[lids[i]]
  
   return lids, gids, x, y, z

def collect_unique_values_over_procs(data, tol=1.0e-10):
   """
   Take a list of scalar data from each processor and return a list of all unique values 
   over all procs to each proc.
   Optional parameter defining the tolerance for scalar comparison
   """

   # let each proc know how many nodes are on the other procs in order to send/receive data
   num_nodes_per_proc = pu.get_all_remote_ints(len(data))
   rank = gt.proc_rank()

   # for each proc, broadcast data to other procs
   # and add any new unique values to the list of all unique values
   unique_values = []
   for p in range(gt.num_procs()):
      # broadcast
      remote_data = np.zeros(num_nodes_per_proc[p], dtype=np.double)
      if rank == p:
        remote_data = data
      pu.broadcast_scalar_data(p, remote_data)
    
      # check for new unique values
      for i in range(num_nodes_per_proc[p]):
        new_value = True
        for j in range(len(unique_values)):
          if abs(remote_data[i] - unique_values[j]) < tol:
            new_value = False
        if new_value:
          unique_values.append(remote_data[i])

   # sort the list of unique values and return
   unique_values.sort()
   return unique_values

def match_coordinates(x, y, ref_x, ref_y, ref_ids, tol=1.0e-10):
   """
   Match a set of coordinates (x,y) against a reference set (ref_x,ref_y) and return
   the corresponding reference IDs for each coordinate pair
   Optional parameter defining the tolerance for scalar comparison
   """

   # let each proc know how many ref nodes are on the other procs in order to send/receive data
   num_ref_nodes_per_proc = pu.get_all_remote_ints(len(ref_ids))
   rank = gt.proc_rank()

   # for each proc, broadcast reference coords to other procs
   # match coordinates against the remote data and assign ids
   rids = -1*np.ones(len(x), dtype=np.longlong)
   for p in range(gt.num_procs()):
      # skip procs with no reference nodes
      if num_ref_nodes_per_proc[p] < 1:
        continue

      # broadcast
      remote_ids = np.zeros(num_ref_nodes_per_proc[p], dtype=np.longlong)
      remote_x   = np.zeros(num_ref_nodes_per_proc[p], dtype=np.double)
      remote_y   = np.zeros(num_ref_nodes_per_proc[p], dtype=np.double)
      if rank == p:
        remote_ids = ref_ids
        remote_x   = ref_x
        remote_y   = ref_y
      pu.broadcast_go_data(p, remote_ids)
      pu.broadcast_scalar_data(p, remote_x)
      pu.broadcast_scalar_data(p, remote_y)

      # match
      for i in range(len(x)):
        # skip nodes that have already been matched
        if rids[i] > -1:
          continue
        for j in range(num_ref_nodes_per_proc[p]):
          if abs(x[i] - remote_x[j]) < tol:
             if abs(y[i] - remote_y[j]) < tol:
                rids[i] = remote_ids[j]

   # error check
   if min(rids) < 0:
      raise Exception('Some nodes have not been assigned IDs corresponding to the reference plane')
   return rids

def distribute_composite_ids_across_procs(num_procs_x,total_x,num_procs_y,total_y):
   """
   Given how many procs to divide a number of ids in two coordinate directions into,
   distribute composite ids on each processor that fit this decomposition
   """

   # assign ranks in each coordinate direction
   rank = gt.proc_rank()
   num_procs = gt.num_procs()
   x_rank = rank*num_procs_x//num_procs
   y_rank = rank%num_procs_y

   # get num xids and start xid on each rank
   num_target_xids = total_x//num_procs_x
   start_xid = x_rank*num_target_xids+min(x_rank,total_x%num_procs_x)
   if x_rank < total_x%num_procs_x:
      num_target_xids = num_target_xids+1

   # get num yids and start yid on each rank
   num_target_yids = total_y//num_procs_y
   start_yid = y_rank*num_target_yids+min(y_rank,total_y%num_procs_y)
   if y_rank < total_y%num_procs_y:
      num_target_yids = num_target_yids+1

   # assign composite ids
   num_target_cids = num_target_yids*num_target_xids
   target_cids = -1*np.ones(num_target_cids, dtype=np.longlong)
   for i in range(num_target_xids):
     for j in range(num_target_yids):
       target_cids[i*num_target_yids+j] = (start_xid+i)*total_y + (start_yid+j)

   # create the global blocking array needed for dist tensor context
   global_blocking = np.zeros((2, num_procs+1), dtype=np.longlong)
   for i in range(num_procs_x):
     global_blocking[0,i+1] = global_blocking[0,i] + total_x//num_procs_x
     if i < total_x%num_procs_x:
       global_blocking[0,i+1] += 1
   for i in range(num_procs_y):
     global_blocking[1,i+1] = global_blocking[1,i] + total_y//num_procs_y
     if i < total_y%num_procs_y:
       global_blocking[1,i+1] += 1

   return target_cids, global_blocking

def distribute_composite_ids_to_root(num_procs_x,total_x,num_procs_y,total_y):
   """
   Given how many procs to divide a number of ids in two coordinate directions into,
   distribute composite ids on each processor that fit this decomposition
   """

   # assign ranks in each coordinate direction
   target_cids = -1*np.ones(0, dtype=np.longlong)
   rank = gt.proc_rank()
   if rank == 0:
      total_cids = total_x*total_y
      target_cids = np.arange(0, total_cids-1, 1, dtype=np.longlong)
   return target_cids

if __name__ == "__main__":
   gt.initializeGenten()
   assert len(sys.argv) >= 3, "usage: torus_to_tensor.py <exodus_base_name> <axis_of_rotation> <num_procs_per_poloidal_plane> <tol>"
   num_procs = gt.num_procs()
   base_filename = sys.argv[1]
   axis = sys.argv[2]
   num_procs_per_poloidal_plane = num_procs
   if len(sys.argv) >= 4:
     num_procs_per_poloidal_plane = int(sys.argv[3])
   tol = 1.0e-10
   if len(sys.argv) >= 5:
     tol = np.double(sys.argv[4])
   tensor, global_blocking, parallel_map, node_gids = torus_to_tensor(base_filename, axis, num_procs_per_poloidal_plane, tol)
   gt.finalizeGenten()


