import base64
import json
import struct
import tensorflow as tf


class SimResult:
  def __init__(self, stream):
    def check(item):
      if type(item) == dict:
        return {k: check(v) for k, v in item.items()}
      elif type(item) == list:
        return list(map(check, item))
      elif type(item) == str and item.startswith('buffer:'):
        return base64.b64decode(item[len('buffer:'):])
      else:
        return item

    self.raw = check(json.load(stream))

    def vertex(obj):
      return [obj['x'], obj['y'], obj['z']]

    self.vertices = tf.convert_to_tensor(list(
        map(vertex, self.raw['model']['bodies']['world']['grid']['vertices'])),
                                         dtype=tf.float32)
    crust = tf.convert_to_tensor(
        list(
            map(
                lambda x: x[0],
                struct.iter_unpack(
                    'f', self.raw['model']['bodies']['world']['lithosphere']
                    ['total_crust']))))
    crust = tf.reshape(crust, (8, self.vertices.shape[0]))
    self.crust_layers = {
        SimResult.crust_columns()[i]: crust[i]
        for i in range(0, 8)
    }
    total_crust = tf.zeros((self.vertices.shape[0], ), dtype=tf.float32)
    for k, v in self.crust_layers.items():
      if k != 'age':
        total_crust += v
    self.crust_layers['total'] = total_crust

  def seed(self):
    return self.raw['seed']

  def crust_columns():
    return [
        'sediment', 'sedimentary', 'metamorphic', 'felsic_plutonic',
        'felsic_volcanic', 'mafic_volcanic', 'mafic_plutonic', 'age'
    ]