import msprime
import numpy as np

import stdpopsim


_species = stdpopsim.get_species("AnoGam")


def GAS_stairwayplot():
    id = "QC-GabonAg1000G_1A17"
    generation_time = 1 / 11

    # Parameters are from Miles et al 2016, which says in Supplement part 8
    # that "All results are available from the Ag1000G FTP site."
    # So, obtained results from:
    # wget ftp://ngs.sanger.ac.uk/production/ag1000g/phase1/\
    # AR3.1/stairway_plot/GAS.meru_mela.sfs.sp.summary
    # tail -n +8 GAS.meru_mela.sfs.sp.summary | cut -f 6,7 > GAS.txt
    # and pasted this into the below

    # (year, Ne_median) pairs:
    # note the repeated time entries
    raw_data = np.array(
        [
            (0.0, 4069863.2641375437),
            (119.04361951964268, 4069863.2641375437),
            (119.04361951964268, 2570469.788804191),
            (195.5969599650061, 2570469.788804191),
            (195.5969599650061, 2398970.4633577284),
            (268.3536651999058, 2398970.4633577284),
            (268.3536651999058, 649423.6408901582),
            (288.41432194750894, 649423.6408901582),
            (288.41432194750894, 647425.872043444),
            (308.78707980783946, 647425.872043444),
            (308.78707980783946, 629386.5533302526),
            (328.96586993080166, 629386.5533302526),
            (328.96586993080166, 623896.5503633062),
            (349.34965006192334, 623896.5503633062),
            (349.34965006192334, 615171.6634657005),
            (369.83488694056797, 615171.6634657005),
            (369.83488694056797, 609770.7169891531),
            (390.53455144559734, 609770.7169891531),
            (390.53455144559734, 609770.7169891531),
            (411.64009172523504, 609770.7169891531),
            (411.64009172523504, 606896.7847736678),
            (433.0621204495311, 606896.7847736678),
            (433.0621204495311, 606896.7847736678),
            (454.912589748313, 606896.7847736678),
            (454.912589748313, 606896.7847736678),
            (477.20448266929253, 606896.7847736678),
            (477.20448266929253, 606896.7847736678),
            (499.9513121804962, 606896.7847736678),
            (499.9513121804962, 606896.7847736678),
            (523.1671484857452, 606896.7847736678),
            (523.1671484857452, 606896.7847736678),
            (546.8666480473536, 606896.7847736678),
            (546.8666480473536, 606896.7847736678),
            (571.065084441838, 606896.7847736678),
            (571.065084441838, 606896.7847736678),
            (595.7783811851413, 606896.7847736678),
            (595.7783811851413, 606896.7847736678),
            (621.0231466756121, 606896.7847736678),
            (621.0231466756121, 606896.7847736678),
            (646.816711415876, 606896.7847736678),
            (646.816711415876, 606896.7847736678),
            (673.1771676888928, 606896.7847736678),
            (673.1771676888928, 606896.7847736678),
            (700.1234118790879, 606896.7847736678),
            (700.1234118790879, 606896.7847736678),
            (727.6751896465906, 606896.7847736678),
            (727.6751896465906, 606896.7847736678),
            (755.8531441815365, 606896.7847736678),
            (755.8531441815365, 609770.7169891531),
            (784.8153706937982, 609770.7169891531),
            (784.8153706937982, 612860.2993041094),
            (814.6012956678746, 612860.2993041094),
            (814.6012956678746, 646628.9792125917),
            (846.7678908283069, 646628.9792125917),
            (846.7678908283069, 672653.8403425539),
            (881.0257894961221, 672653.8403425539),
            (881.0257894961221, 694949.5701181502),
            (917.272050911558, 694949.5701181502),
            (917.272050911558, 2053484.489041333),
            (1026.987248146155, 2053484.489041333),
            (1026.987248146155, 2512451.153376244),
            (1164.5389797372006, 2512451.153376244),
            (1164.5389797372006, 2726859.95574373),
            (1317.561423942797, 2726859.95574373),
            (1317.561423942797, 2779609.283863392),
            (1477.4929132329808, 2779609.283863392),
            (1477.4929132329808, 2827486.5910723945),
            (1644.3505799334257, 2827486.5910723945),
            (1644.3505799334257, 2833326.9048009124),
            (1815.8958168056997, 2833326.9048009124),
            (1815.8958168056997, 2830283.032594962),
            (1991.7662593953326, 2830283.032594962),
            (1991.7662593953326, 2726859.95574373),
            (2165.728617018537, 2726859.95574373),
            (2165.728617018537, 2707367.4282596707),
            (2343.115508305002, 2707367.4282596707),
            (2343.115508305002, 2557240.922090161),
            (2515.2565282700093, 2557240.922090161),
            (2515.2565282700093, 2420531.1968794307),
            (2682.7209808007574, 2420531.1968794307),
            (2682.7209808007574, 2245241.112841403),
            (2842.433679110716, 2245241.112841403),
            (2842.433679110716, 1829782.2297892452),
            (2976.3120204159472, 1829782.2297892452),
            (2976.3120204159472, 743769.2583260541),
            (3032.3082002259275, 743769.2583260541),
            (3032.3082002259275, 719107.7911294184),
            (3088.040029345467, 719107.7911294184),
            (3088.040029345467, 710882.4752727259),
            (3144.778992971672, 710882.4752727259),
            (3144.778992971672, 708782.9061540969),
            (3203.064664297201, 708782.9061540969),
            (3203.064664297201, 706605.8391406796),
            (3262.9592045930835, 706605.8391406796),
            (3262.9592045930835, 702884.8367882286),
            (3324.400186829817, 702884.8367882286),
            (3324.400186829817, 699330.1487743513),
            (3387.4710876933045, 699330.1487743513),
            (3387.4710876933045, 698106.7921761848),
            (3452.462643079068, 698106.7921761848),
            (3452.462643079068, 697332.0358451155),
            (3519.5105769135102, 697332.0358451155),
            (3519.5105769135102, 696317.9346790552),
            (3588.692686668012, 696317.9346790552),
            (3588.692686668012, 696167.6446096608),
            (3660.2045145683364, 696167.6446096608),
            (3660.2045145683364, 695633.5594190623),
            (3734.1255133908285, 695633.5594190623),
            (3734.1255133908285, 695633.5594190623),
            (3810.64023147025, 695633.5594190623),
            (3810.64023147025, 695633.5594190623),
            (3889.8876180525085, 695633.5594190623),
            (3889.8876180525085, 695633.5594190623),
            (3972.0167277832124, 695633.5594190623),
            (3972.0167277832124, 695633.5594190623),
            (4057.1876563928317, 695633.5594190623),
            (4057.1876563928317, 695633.5594190623),
            (4145.572582308474, 695633.5594190623),
            (4145.572582308474, 695633.5594190623),
            (4237.356928451641, 695633.5594190623),
            (4237.356928451641, 695633.5594190623),
            (4332.740660718069, 695633.5594190623),
            (4332.740660718069, 695633.5594190623),
            (4431.939742275155, 695633.5594190623),
            (4431.939742275155, 695633.5594190623),
            (4535.187765936612, 695633.5594190623),
            (4535.187765936612, 695633.5594190623),
            (4642.737790583962, 695633.5594190623),
            (4642.737790583962, 695633.5594190623),
            (4754.864412024816, 695633.5594190623),
            (4754.864412024816, 695633.5594190623),
            (4871.8661039631, 695633.5594190623),
            (4871.8661039631, 695633.5594190623),
            (4994.06787109864, 695633.5594190623),
            (4994.06787109864, 695633.5594190623),
            (5121.824264013069, 695633.5594190623),
            (5121.824264013069, 695633.5594190623),
            (5255.52281473747, 695633.5594190623),
            (5255.52281473747, 694677.6052659544),
            (5395.395482665781, 694677.6052659544),
            (5395.395482665781, 694677.6052659544),
            (5542.0912075662045, 694677.6052659544),
            (5542.0912075662045, 694677.6052659544),
            (5696.121718711648, 694677.6052659544),
            (5696.121718711648, 694677.6052659544),
            (5858.051230428654, 694677.6052659544),
            (5858.051230428654, 694677.6052659544),
            (6028.5033480255015, 694677.6052659544),
            (6028.5033480255015, 694677.6052659544),
            (6208.169093600558, 694677.6052659544),
            (6208.169093600558, 693429.2426962901),
            (6397.475466163017, 693429.2426962901),
            (6397.475466163017, 693429.2426962901),
            (6597.599345729046, 693429.2426962901),
            (6597.599345729046, 693429.2426962901),
            (6809.4952182107245, 693429.2426962901),
            (6809.4952182107245, 692716.935261925),
            (7034.002408378257, 692716.935261925),
            (7034.002408378257, 692379.4193946248),
            (7272.425073459051, 692379.4193946248),
            (7272.425073459051, 690176.6983717981),
            (7525.4223969091245, 690176.6983717981),
            (7525.4223969091245, 688681.8347837181),
            (7794.7017066974795, 688681.8347837181),
            (7794.7017066974795, 687162.8678362874),
            (8081.917116451414, 687162.8678362874),
            (8081.917116451414, 686235.3636992247),
            (8389.232550262084, 686235.3636992247),
            (8389.232550262084, 686150.9091049897),
            (8719.271467580502, 686150.9091049897),
            (8719.271467580502, 685951.341450253),
            (9074.594617774881, 685951.341450253),
            (9074.594617774881, 685005.950816014),
            (9457.814730119504, 685005.950816014),
            (9457.814730119504, 684576.5933044655),
            (9872.709635152514, 684576.5933044655),
            (9872.709635152514, 684576.5933044655),
            (10323.682358014481, 684576.5933044655),
            (10323.682358014481, 683889.4362766537),
            (10815.15877646698, 683889.4362766537),
            (10815.15877646698, 683889.4362766537),
            (11353.442472867337, 683889.4362766537),
            (11353.442472867337, 683889.4362766537),
            (11945.554538907729, 683889.4362766537),
            (11945.554538907729, 683889.4362766537),
            (12599.994190847112, 683889.4362766537),
            (12599.994190847112, 687172.5863881508),
            (13330.640215714893, 687172.5863881508),
            (13330.640215714893, 1980227.2156260633),
            (15683.851404813011, 1980227.2156260633),
            (15683.851404813011, 1989715.587026168),
            (18343.898981051203, 1989715.587026168),
            (18343.898981051203, 1992430.3521081463),
            (21362.732847881725, 1992430.3521081463),
            (21362.732847881725, 1998953.4405011814),
            (24824.12408684481, 1998953.4405011814),
            (24824.12408684481, 1998953.4405011814),
            (28818.037054879136, 1998953.4405011814),
            (28818.037054879136, 1998160.881282157),
            (33475.75472686552, 1998160.881282157),
            (33475.75472686552, 1998160.881282157),
            (38980.33015739488, 1998160.881282157),
            (38980.33015739488, 1998160.881282157),
            (45585.8206740301, 1998160.881282157),
            (45585.8206740301, 1972764.0343926845),
            (53556.584449354086, 1972764.0343926845),
            (53556.584449354086, 1903551.7604157536),
            (63170.48222923164, 1903551.7604157536),
            (63170.48222923164, 543762.2932593989),
            (66701.40621143552, 543762.2932593989),
            (66701.40621143552, 530289.214160486),
            (71292.65481888562, 530289.214160486),
            (71292.65481888562, 624896.6793768455),
            (78867.16002345344, 624896.6793768455),
            (78867.16002345344, 1680052.006021),
            (109413.56013292617, 1680052.006021),
            (109413.56013292617, 322144.5653929737),
            (119175.51665998598, 322144.5653929737),
            (119175.51665998598, 77334.77862527313),
            (123862.47294030557, 77334.77862527313),
            (123862.47294030557, 409527.1280141645),
            (198321.95076106273, 409527.1280141645),
        ]
    )

    model = msprime.Demography()
    model.add_population(
        name="GAS", description="Gabon sample", initial_size=raw_data[0][1]
    )
    last_time, last_size = 0, raw_data[0][1]
    for j in range(1, raw_data.shape[0]):
        time, size = raw_data[j, :]
        if time > last_time and last_size != size:
            model.add_population_parameters_change(
                time * generation_time,
                initial_size=size,
                growth_rate=0,
                population="GAS",
            )
            last_time, last_size = time, size

    return stdpopsim.DemographicModel(
        id=id,
        description=id,
        long_description=id,
        generation_time=generation_time,
        mutation_rate=3.5e-9,
        model=model,
    )


_species.get_demographic_model("GabonAg1000G_1A17").register_qc(GAS_stairwayplot())
