import matplotlib.pyplot as plt
import numpy as np
from numpy import array
from moments import Misc, Spectrum, Numerics, Manips, Integration, Demographics1D, Demographics2D
import sys, os
import argparse
os.chdir("/home/cbattey2/speciation_cyclical_migration/")
parser = argparse.ArgumentParser()
parser.add_argument("--infile")
args = parser.parse_args()
#args=argparse.Namespace(infile="simulations/msprime/constant/constant_sim_10441991.trees.vcf")
# import data
print(args.infile)
dd = Misc.make_data_dict_vcf(vcf_filename=args.infile,
popinfo_filename="id_pop_moments.txt")
pop_ids = ["1", "2"]
projections = [40, 40]
data = Spectrum.from_data_dict(dd, pop_ids, projections, polarized=False)
ns = data.sample_sizes
np.set_printoptions(precision=3)
# print info about sfs
print("projection", projections)
print("sample sizes", data.sample_sizes)
sfs_sum = np.around(data.S(), 2)
print("Sum of SFS = ", sfs_sum, '\n', '\n')
# define model
def IM(params, ns):
