# Success Criteria for this exercise is a subplot containing 2 Heatmaps side by side in the same Figure
from pdb import set_trace
from covidDash.plot_handlers.colorscales import bone_r
from plotly.subplots import make_subplots
import plotly.figure_factory as ff
import plotly.graph_objects as go


# DATA PREP SECTION
# wow Heatmap data
wow = {'x_labels' : {'name' : 'Week',
                     'labels' : ['W21', 'W22', 'W23', 'W24', 'W25', 'W26']
                    },
       'y_labels' : {'name' : 'Site',
                     'labels' : ['A', 'B', 'C', 'D', 'E', 'F', 'G']
                    },
       'data'     : [
                      [0, 0, 1, 0, 0, 0],
                      [0, 0, 3, 1, 0, 0],
                      [0, 0, 0, 0, 0, 0],
                      [0, 0, 0, 0, 0, 1],
                      [0, 0, 0, 0, 0, 0],
                      [0, 0, 1, 0, 0, 0],
                      [0, 0, 0, 0, 0, 0]
                    ],
       'data_labels' : []
      }


# 6w and YTD Heatmap data
totals = {'x_labels' : {'name' : 'Week',
                        'labels' :['6W', 'YTD' ]
         },
          'y_labels' : wow['y_labels'],
          'data'     : [
                         [1, 16],
                         [4, 8],
                         [0, 1],
                         [1, 12],
                         [0, 5],
                         [1, 17],
                         [0, 1]
                      ],
         'data_labels' : []
         }


# this function is simply a base func for now
def int_to_str(arr2d):
  """base function for handling data to label conversion
  Args:

    arr2d (list): a 2D array with numeric values

  Returns:

    r_data (list): a 2D array with values converted into strings

  """
  r_data = []
  for row in arr2d:
    new_row = []
    [new_row.append(str(n)) for n in row]
    r_data.append(new_row)
  return r_data


wow['data_labels'] = int_to_str(wow['data'])
totals['data_labels'] = int_to_str(totals['data'])

# PLOT PREP SECTION
# colorbar placement
wow_cbar= {
      'x' : 1.0,
      'title' : {
        'text' : 'Wow',
        'side' : 'right'
      }
    }


total_cbar = {
        'x' : 1.05,
        'title' : {
          'text' : 'Totals',
          'side' : 'right'
         }
       }

# xaxis conf
xaxis_conf={'rangeslider': {'visible': True},
       'type' : 'category',
       'side' : 'top'
      }


# week over week figure
fig1 = ff.create_annotated_heatmap(x=wow['x_labels']['labels'],
                                  y=wow['y_labels']['labels'],
                                  z=wow['data'],
                                  colorscale=bone_r,
                                  font_colors=['black','white'],
                                  showscale=True,
                                  annotation_text=wow['data_labels'],
                                  colorbar=wow_cbar,
                                  name='Wow'
                                  )


# 6W and YTD
fig2 =ff.create_annotated_heatmap(x=totals['x_labels']['labels'],
                                  y=totals['y_labels']['labels'],
                                  z=totals['data'],
                                  colorscale=bone_r,
                                  font_colors=['black','white'],
                                  showscale=True,
                                  annotation_text=totals['data_labels'],
                                  colorbar=total_cbar,
                                  name='Totals',
                                  )
# SUBPLOT PREP SECTION
# base subplot
fig = make_subplots(
    rows=1, cols=2,
    shared_yaxes=True,
    horizontal_spacing=0,
    row_titles=wow['y_labels']['labels'],
#    y_title=[wow['y_labels']['name'],
    x_title=wow['x_labels']['name'],
    column_widths=[0.75, 0.25]
)

# add data
fig.add_trace(fig1.data[0], 1, 1)
fig.add_trace(fig2.data[0], 1, 2)

# apply annotations
wow_annot = list(fig1.layout.annotations)
totals_annot = list(fig2.layout.annotations)
for k in range(len(totals_annot)):
  totals_annot[k]['xref'] = 'x2'
  totals_annot[k]['yref'] = 'y2'

fig.update_layout(annotations=wow_annot+totals_annot,xaxis=xaxis_conf, xaxis2={'side':'top'})
set_trace()

fig.show()