import numpy as np
np.random.seed(1)
import scipy.stats
import random
random.seed(1)
import matplotlib.pyplot as plt
%matplotlib inline
import sys
# import anhima
# dev imports
sys.path.insert(0, '..')
%reload_ext autoreload
%autoreload 1
%aimport anhima.sim
%aimport anhima.gt
%aimport anhima.dist

# simulate genotypes with no relatedness
n_variants = 1000
n_samples = 100
ploidy = 2
af_dist = scipy.stats.beta(a=.4, b=.6)
p_missing = 0
genotypes = anhima.sim.simulate_biallelic_genotypes(n_variants, n_samples, af_dist, p_missing, ploidy)
gn = anhima.gt.as_n_alt(genotypes)

# simulate three sub-populations with relatedness
genotypes_subpop1 = anhima.sim.simulate_relatedness(genotypes[:, :n_samples//2], relatedness=.5, n_iter=1000)
genotypes_subpop2 = anhima.sim.simulate_relatedness(genotypes[:, n_samples//2:], relatedness=.5, n_iter=1000)
genotypes_subpop2a = anhima.sim.simulate_relatedness(genotypes_subpop2[:, n_samples//4:], relatedness=.5, n_iter=500)
genotypes_subpop2b = anhima.sim.simulate_relatedness(genotypes_subpop2[:, :n_samples//4], relatedness=.5, n_iter=500)
genotypes_related = np.concatenate([genotypes_subpop1, genotypes_subpop2a, genotypes_subpop2b], axis=1)
gnr = anhima.gt.as_n_alt(genotypes_related)

# calculate pairwise distance (Euclidean is default metric)
d, ds = anhima.dist.pairwise_distance(gn)
dr, dsr = anhima.dist.pairwise_distance(gnr)

plt.hist(d, bins=30);

plt.hist(dr, bins=30);

anhima.dist.plot_pairwise_distance(ds);

anhima.dist.plot_pairwise_distance(dsr);

d, s = anhima.dist.pairwise_distance(gnr, metric='euclidean')
plt.hist(d, bins=30)
plt.xlim(min(d), max(d));

d, s = anhima.dist.pairwise_distance(gnr, metric='sqeuclidean')
plt.hist(d, bins=30)
plt.xlim(min(d), max(d));

d, s = anhima.dist.pairwise_distance(gnr, metric='cityblock')
plt.hist(d, bins=30)
plt.xlim(min(d), max(d));