import os

import mock
import pytest

from program_synthesis.karel import arguments
from program_synthesis.karel.dataset import dataset
from program_synthesis.karel.dataset import refine_env
from program_synthesis.karel.models import karel_edit_model


@pytest.fixture
def args():
    with mock.patch('sys.argv', [
            'test', '--num_placeholders', '0', '--karel-merge-io', 'setlstm'
    ]):
        args = arguments.parse('', 'train')
        args.word_vocab = os.path.abspath(
            os.path.join(
                os.path.dirname(os.path.realpath(__file__)),
                '../../data/karel/word.vocab'))
        return args


def train_two_steps(args, current_codes):
    assert len(current_codes) == len(examples)

    batch = []
    for example, current_code in zip(examples, current_codes):
        goal_atree = refine_env.AnnotatedTree(code=example['code'])
        current_atree = refine_env.AnnotatedTree(code=current_code)

        actions = refine_env.ComputeAddOps.run(current_atree, goal_atree)

        batch.append(
            dataset.KarelEditExample(current_code, example['code'], actions,
                                     example['input_tests'][:5], example[
                                         'input_tests'][5:]))

    m = karel_edit_model.KarelStepEditModel(args)
    batch_processor = m.batch_processor(for_eval=False)

    loss1 = m.train(batch_processor(batch))['loss']
    loss2 = m.train(batch_processor(batch))['loss']

    assert loss2 < loss1


def test_train_from_empty(args):
    train_two_steps(args, ['DEF run m( m)'.split(), 'DEF run m( m)'.split()])


def test_train_from_nearly_full(args):
    train_two_steps(args, [
        'DEF run m( REPEAT R=5 r( turnLeft turnRight move r) pickMarker m)'.
        split(),
        'DEF run m( turnLeft REPEAT R=2 r( pickMarker r) m)'.split(),
    ])


def test_train_from_full(args):
    train_two_steps(args, [ex['code'] for ex in examples])


examples = [
    # Index 0 in train.pkl
    {
        'code': 'DEF run m( REPEAT R=5 r( turnLeft IFELSE c( markersPresent c)'
        ' i( turnRight i) ELSE e( move e) r) pickMarker m)'.split(),
        'input_tests': [{
            'input': [
                811, 1325, 1360, 1379, 1405, 1411, 1423, 1431, 1472, 1480,
                1620, 1621, 1622, 1623, 1624, 1625, 1626, 1627, 1628, 1629,
                1630, 1631, 1632, 1633, 1634, 1635, 1638, 1653, 1656, 1671,
                1674, 1689, 1692, 1707, 1710, 1725, 1728, 1743, 1746, 1761,
                1764, 1779, 1782, 1797, 1800, 1815, 1818, 1819, 1820, 1821,
                1822, 1823, 1824, 1825, 1826, 1827, 1828, 1829, 1830, 1831,
                1832, 1833, 1976, 1999, 2076, 2095, 2096, 2099, 2108, 2117,
                2130, 2635, 2636, 2766, 3016
            ],
            'output': [
                488, 1325, 1360, 1379, 1405, 1411, 1423, 1431, 1472, 1480,
                1620, 1621, 1622, 1623, 1624, 1625, 1626, 1627, 1628, 1629,
                1630, 1631, 1632, 1633, 1634, 1635, 1638, 1653, 1656, 1671,
                1674, 1689, 1692, 1707, 1710, 1725, 1728, 1743, 1746, 1761,
                1764, 1779, 1782, 1797, 1800, 1815, 1818, 1819, 1820, 1821,
                1822, 1823, 1824, 1825, 1826, 1827, 1828, 1829, 1830, 1831,
                1832, 1833, 1976, 1999, 2076, 2095, 2096, 2099, 2117, 2130,
                2635, 2636, 2766, 3016
            ]
        }, {
            'input': [
                1159, 1320, 1333, 1334, 1358, 1405, 1423, 1620, 1621, 1622,
                1623, 1624, 1625, 1626, 1627, 1628, 1629, 1630, 1638, 1648,
                1656, 1666, 1674, 1684, 1692, 1702, 1710, 1720, 1728, 1738,
                1746, 1756, 1764, 1774, 1782, 1792, 1800, 1810, 1818, 1819,
                1820, 1821, 1822, 1823, 1824, 1825, 1826, 1827, 1828, 1965,
                2043, 2077, 2089, 2093, 2113, 2453, 2775, 3267, 3979, 4364,
                4599, 4704
            ],
            'output': [
                817, 1320, 1333, 1334, 1358, 1405, 1423, 1620, 1621, 1622,
                1623, 1624, 1625, 1626, 1627, 1628, 1629, 1630, 1638, 1648,
                1656, 1666, 1674, 1684, 1692, 1702, 1710, 1720, 1728, 1738,
                1746, 1756, 1764, 1774, 1782, 1792, 1800, 1810, 1818, 1819,
                1820, 1821, 1822, 1823, 1824, 1825, 1826, 1827, 1828, 1965,
                2043, 2077, 2089, 2093, 2453, 2775, 3267, 3979, 4364, 4599,
                4704
            ]
        }, {
            'input': [
                1123, 1317, 1333, 1334, 1335, 1358, 1371, 1375, 1376, 1391,
                1620, 1621, 1622, 1623, 1624, 1625, 1626, 1627, 1628, 1629,
                1638, 1647, 1656, 1665, 1674, 1683, 1692, 1701, 1710, 1719,
                1728, 1737, 1746, 1755, 1764, 1773, 1782, 1791, 1800, 1801,
                1802, 1803, 1804, 1805, 1806, 1807, 1808, 1809, 2077, 2093,
                2109, 2110, 2974, 4560, 4706
            ],
            'output': [
                781, 1317, 1333, 1334, 1335, 1358, 1371, 1375, 1376, 1391,
                1620, 1621, 1622, 1623, 1624, 1625, 1626, 1627, 1628, 1629,
                1638, 1647, 1656, 1665, 1674, 1683, 1692, 1701, 1710, 1719,
                1728, 1737, 1746, 1755, 1764, 1773, 1782, 1791, 1800, 1801,
                1802, 1803, 1804, 1805, 1806, 1807, 1808, 1809, 2093, 2109,
                2110, 2974, 4560, 4706
            ]
        }, {
            'input': [
                1100, 1352, 1620, 1621, 1622, 1623, 1638, 1641, 1656, 1659,
                1674, 1677, 1692, 1695, 1710, 1713, 1728, 1731, 1746, 1749,
                1764, 1765, 1766, 1767, 2054
            ],
            'output': [
                758, 1352, 1620, 1621, 1622, 1623, 1638, 1641, 1656, 1659,
                1674, 1677, 1692, 1695, 1710, 1713, 1728, 1731, 1746, 1749,
                1764, 1765, 1766, 1767
            ]
        }, {
            'input': [
                759, 1317, 1335, 1337, 1351, 1406, 1441, 1448, 1620, 1621,
                1622, 1623, 1624, 1625, 1626, 1627, 1628, 1629, 1630, 1638,
                1648, 1656, 1666, 1674, 1684, 1692, 1702, 1710, 1720, 1728,
                1738, 1746, 1756, 1764, 1774, 1782, 1792, 1800, 1801, 1802,
                1803, 1804, 1805, 1806, 1807, 1808, 1809, 1810, 2022, 2023,
                2056, 2291, 3030, 3584
            ],
            'output': [
                436, 1317, 1335, 1337, 1351, 1406, 1441, 1448, 1620, 1621,
                1622, 1623, 1624, 1625, 1626, 1627, 1628, 1629, 1630, 1638,
                1648, 1656, 1666, 1674, 1684, 1692, 1702, 1710, 1720, 1728,
                1738, 1746, 1756, 1764, 1774, 1782, 1792, 1800, 1801, 1802,
                1803, 1804, 1805, 1806, 1807, 1808, 1809, 1810, 2022, 2023,
                2291, 3030, 3584
            ]
        }]
    },
    # Index 2 in train.pkl
    {
        'code':
        'DEF run m( turnLeft REPEAT R=2 r( move pickMarker r) m)'.split(),
        'input_tests': [{
            'input': [
                361, 1388, 1389, 1620, 1621, 1622, 1623, 1624, 1625, 1638,
                1643, 1656, 1661, 1674, 1679, 1692, 1697, 1710, 1715, 1728,
                1729, 1730, 1731, 1732, 1733, 1965, 1981, 1982, 2017, 2035,
                2287, 2971, 4251
            ],
            'output': [
                73, 1388, 1389, 1620, 1621, 1622, 1623, 1624, 1625, 1638, 1643,
                1656, 1661, 1674, 1679, 1692, 1697, 1710, 1715, 1728, 1729,
                1730, 1731, 1732, 1733, 1965, 1981, 1982, 2035, 2287, 2647,
                4251
            ]
        }, {
            'input': [
                434, 1620, 1621, 1622, 1623, 1624, 1625, 1626, 1627, 1628,
                1629, 1630, 1631, 1638, 1649, 1656, 1667, 1674, 1685, 1692,
                1703, 1710, 1721, 1728, 1739, 1746, 1757, 1764, 1775, 1782,
                1783, 1784, 1785, 1786, 1787, 1788, 1789, 1790, 1791, 1792,
                1793, 1964, 1999, 2005, 2021, 2024, 2039, 2060, 2072, 3062,
                3066
            ],
            'output': [
                146, 1620, 1621, 1622, 1623, 1624, 1625, 1626, 1627, 1628,
                1629, 1630, 1631, 1638, 1649, 1656, 1667, 1674, 1685, 1692,
                1703, 1710, 1721, 1728, 1739, 1746, 1757, 1764, 1775, 1782,
                1783, 1784, 1785, 1786, 1787, 1788, 1789, 1790, 1791, 1792,
                1793, 1964, 1999, 2005, 2021, 2024, 2039, 2060, 2738, 3066
            ]
        }, {
            'input': [
                343, 1620, 1621, 1622, 1623, 1624, 1625, 1626, 1627, 1628,
                1629, 1630, 1631, 1632, 1633, 1638, 1651, 1656, 1669, 1674,
                1687, 1692, 1705, 1710, 1723, 1728, 1741, 1746, 1759, 1764,
                1777, 1782, 1795, 1800, 1813, 1818, 1831, 1836, 1849, 1854,
                1867, 1872, 1885, 1890, 1903, 1908, 1909, 1910, 1911, 1912,
                1913, 1914, 1915, 1916, 1917, 1918, 1919, 1920, 1921, 1967,
                1981, 2005, 2038, 2058, 2077, 2079, 2125, 2128, 2132, 2136,
                2218, 2647, 3479, 3584, 3997, 4323, 4566
            ],
            'output': [
                55, 1620, 1621, 1622, 1623, 1624, 1625, 1626, 1627, 1628, 1629,
                1630, 1631, 1632, 1633, 1638, 1651, 1656, 1669, 1674, 1687,
                1692, 1705, 1710, 1723, 1728, 1741, 1746, 1759, 1764, 1777,
                1782, 1795, 1800, 1813, 1818, 1831, 1836, 1849, 1854, 1867,
                1872, 1885, 1890, 1903, 1908, 1909, 1910, 1911, 1912, 1913,
                1914, 1915, 1916, 1917, 1918, 1919, 1920, 1921, 1967, 2005,
                2038, 2058, 2077, 2079, 2125, 2128, 2132, 2136, 2218, 2323,
                3479, 3584, 3997, 4323, 4566
            ]
        }, {
            'input': [
                542, 1374, 1390, 1407, 1409, 1425, 1441, 1442, 1463, 1477,
                1482, 1500, 1513, 1572, 1620, 1621, 1622, 1623, 1624, 1625,
                1626, 1627, 1638, 1645, 1656, 1663, 1674, 1681, 1692, 1699,
                1710, 1717, 1728, 1735, 1746, 1753, 1764, 1771, 1782, 1789,
                1800, 1807, 1818, 1825, 1836, 1843, 1854, 1861, 1872, 1879,
                1890, 1897, 1908, 1909, 1910, 1911, 1912, 1913, 1914, 1915,
                1963, 1965, 1967, 1982, 1984, 1986, 1999, 2002, 2019, 2020,
                2040, 2053, 2056, 2058, 2071, 2072, 2075, 2076, 2092, 2094,
                2107, 2108, 2109, 2112, 2126, 2127, 2129, 2145, 2146, 2147,
                2162, 2163, 2165, 2166, 2180, 2182, 2184, 2197, 2200, 2215,
                2216, 2218, 2292, 2307, 2342, 2417, 2543, 2648, 2687, 2702,
                2865, 2993, 3082, 3151, 3155, 3170, 3171, 3281, 3333, 3439,
                3655, 3711, 3822, 4072, 4088, 4108, 4249, 4556, 4595, 4609,
                4793
            ],
            'output': [
                254, 1374, 1390, 1407, 1409, 1425, 1441, 1442, 1463, 1477,
                1482, 1500, 1513, 1572, 1620, 1621, 1622, 1623, 1624, 1625,
                1626, 1627, 1638, 1645, 1656, 1663, 1674, 1681, 1692, 1699,
                1710, 1717, 1728, 1735, 1746, 1753, 1764, 1771, 1782, 1789,
                1800, 1807, 1818, 1825, 1836, 1843, 1854, 1861, 1872, 1879,
                1890, 1897, 1908, 1909, 1910, 1911, 1912, 1913, 1914, 1915,
                1963, 1965, 1967, 1982, 1984, 1986, 1999, 2002, 2019, 2020,
                2040, 2053, 2056, 2058, 2071, 2072, 2075, 2076, 2092, 2094,
                2107, 2108, 2109, 2112, 2126, 2127, 2129, 2145, 2146, 2147,
                2162, 2163, 2165, 2166, 2182, 2184, 2197, 2200, 2215, 2216,
                2218, 2292, 2307, 2342, 2417, 2543, 2648, 2687, 2702, 2846,
                2865, 2993, 3082, 3151, 3155, 3171, 3281, 3333, 3439, 3655,
                3711, 3822, 4072, 4088, 4108, 4249, 4556, 4595, 4609, 4793
            ]
        }, {
            'input': [
                689, 1319, 1336, 1394, 1406, 1409, 1412, 1620, 1621, 1622,
                1623, 1624, 1625, 1626, 1627, 1628, 1629, 1638, 1647, 1656,
                1665, 1674, 1683, 1692, 1701, 1710, 1719, 1728, 1737, 1746,
                1747, 1748, 1749, 1750, 1751, 1752, 1753, 1754, 1755, 1965,
                1982, 1985, 1987, 1988, 2005, 2036, 2038, 2039, 2058, 2292,
                2293, 2344, 2361, 2958, 2972, 3320, 4555
            ],
            'output': [
                367, 1319, 1336, 1394, 1406, 1409, 1412, 1620, 1621, 1622,
                1623, 1624, 1625, 1626, 1627, 1628, 1629, 1638, 1647, 1656,
                1665, 1674, 1683, 1692, 1701, 1710, 1719, 1728, 1737, 1746,
                1747, 1748, 1749, 1750, 1751, 1752, 1753, 1754, 1755, 1965,
                1982, 1985, 1988, 2005, 2036, 2038, 2039, 2058, 2292, 2293,
                2344, 2361, 2634, 2972, 3320, 4555
            ]
        }]
    }
]
