import math
import numpy as np

def correl(x, y):
  tmp = np.cov(x, y)
  return tmp[0][1] / math.sqrt(tmp[0][0]*tmp[1][1])

def makeGametes(parent_genes):
  # Randomly choose one chromosome from each chromsome-pair.
  n = parent_genes.shape[0]
  chromosome_pairs = parent_genes.shape[1]
  i = np.arange(0, n * chromosome_pairs * 2, 2) + np.random.randint(0, 2, n * chromosome_pairs)
  return parent_genes.flatten()[i].reshape((n, chromosome_pairs))

def sim(n, chromosome_pairs, h, mating_r, estimate_num):
  # This model assumes 0 common environment.
  # It assumes the correlation between mom and dad phenotype is driven by equal correlations between mom and dad phenotype and mom and dad unshared environment.

  # Construct a random dad.
  dad_genes = np.random.normal(0, 1, (n, chromosome_pairs, 2))
  # dad_genotypes = np.sum(dad_genes, axis=(1,2)) / math.sqrt(2 * chromosome_pairs)
  # dad_env = np.random.normal(0, 1, n)
  # dad_phenotypes = h * dad_genotypes + math.sqrt(1 - h*h) * dad_env

  # Construct a mom such that her phenotype correlates with the dad's phenotype with the given `mating_r`.
  mom_genes = mating_r * dad_genes + np.random.normal(0, math.sqrt(1 - mating_r*mating_r), (n, chromosome_pairs, 2))
  # mom_genotypes = np.sum(mom_genes, axis=(1,2)) / math.sqrt(2 * chromosome_pairs)
  # mom_env = mating_r * dad_env + np.random.normal(0, math.sqrt(1-mating_r*mating_r), n)
  # mom_phenotypes = h * mom_genotypes + math.sqrt(1 - h*h) * mom_env

  # Construct child and monozygotic twin.
  sperm = makeGametes(dad_genes)
  eggs = makeGametes(mom_genes)
  child_genotypes = (np.sum(sperm, axis=1) + np.sum(eggs, axis=1)) / math.sqrt(2 * chromosome_pairs)
  child_phenotypes = h * child_genotypes + np.random.normal(0, math.sqrt(1 - h*h), n)
  mz_phenotypes = h * child_genotypes + np.random.normal(0, math.sqrt(1 - h*h), n)

  # Construct dizygotic twin.
  sperm = makeGametes(dad_genes)
  eggs = makeGametes(mom_genes)
  dz_genotypes = (np.sum(sperm, axis=1) + np.sum(eggs, axis=1)) / math.sqrt(2 * chromosome_pairs)
  dz_phenotypes = h * dz_genotypes + np.random.normal(0, math.sqrt(1 - h*h), n)

  # Compute and return what a naive twin study would yield as an estimate of heritability.
  r_mz = correl(child_phenotypes, mz_phenotypes)
  r_dz = correl(child_phenotypes, dz_phenotypes)
  if estimate_num == 0:
    return 2*(r_mz - r_dz)
  elif estimate_num == 1:
    return r_mz
  else:
    assert False

print('<table><thead>')
print('  <tr><td>h<sup>2</sup><td colspan="11" style="text-align: center;">Father-Mother Correlation</td></tr>')
print('  <tr><td></td>' + ''.join(f'<td>{(x/10):.2f}</td>' for x in range(11)) + '</tr>')
print('</thead><tbody>')
for i in range(11):
  h2 = i / 10
  html = '  <tr><td>' + str(round(100*h2)) + '%</td>'
  for j in range(11):
    r = j / 10
    np.random.seed(51243)
    x = sim(1000000, 23, math.sqrt(h2), r, 0)
    html += '<td>' + f"{x:.3f}" + '</td>'
  html += '</tr>'
  print(html)
print('</tbody></table>')