from evalplus.sanitize import code_extract, sanitize


def test_code_extract():
    test_simple = r"""Here is some python code generated
import numpy as np
Sorry, I made a mistake, let me try again
from numpy import sin, cos, tan

def f(x):
    return tan(x)
As you can observe from above
    """
    assert (
        code_extract(test_simple)
        == r"""from numpy import sin, cos, tan

def f(x):
    return tan(x)"""
    )

    test_empty_lines = r"""import numpy as np


import pandas
Sorry, let me try again
from numpy import sin, cos, tan
def f(x):
    return tan(x)
"""
    assert (
        code_extract(test_empty_lines)
        == r"""from numpy import sin, cos, tan
def f(x):
    return tan(x)"""
    )


def test_sanitize_simple():
    icode = r"""Following is the code snippet:
```python
import numpy as np
from numpy import sin, cos

def f(x):
    return np.tan(x)

def g(x):
    return cos(f(x))

def g(x):
    return sin(f(x))

def c(x):
    assert 1==1

assert g(0) == 1
```
"""
    assert (
        sanitize(icode)
        == r"""import numpy as np
from numpy import sin, cos
def f(x):
    return np.tan(x)
def g(x):
    return cos(f(x))"""
    )


def test_sanitize_class():
    icode = r"""Following is the code snippet:
```python
import numpy as np
from numpy import sin, cos
class g():
    def hello_world():
        return 0
def f(x):
    print(g.hello_world())
    return np.tan(x)
```
"""

    assert (
        sanitize(icode)
        == r"""import numpy as np
from numpy import sin, cos
class g():
    def hello_world():
        return 0
def f(x):
    print(g.hello_world())
    return np.tan(x)"""
    )


def test_entrypoint_basic():
    icode = r"""Following is the code snippet:
```python
import numpy as np
from numpy import sin, cos

def f(x):
    return np.tan(x)

def g(x):
    return cos(f(x))

def g(x):
    return sin(f(x))

def c(x):
    return 0

assert g(0) == 1
```
"""
    assert (
        sanitize(icode, "g")
        == r"""import numpy as np
from numpy import sin, cos
def f(x):
    return np.tan(x)
def g(x):
    return cos(f(x))"""
    )


def test_entrypoint_chain():
    icode = r"""Following is the code snippet:
```python
import numpy as np
from numpy import sin, cos

def f(x):
    return c(x)
assert f(1) == 5
def g(x):
    return cos(f(x))

def c(x):
    newObj = h()
    return x

class h():
    def hello_world():
        return 0

class h():
    def goodbye_world():
        return 0


assert g(0) == 1
```
"""
    print(sanitize(icode, "g"))
    assert (
        sanitize(icode, "g")
        == r"""import numpy as np
from numpy import sin, cos
def f(x):
    return c(x)
def g(x):
    return cos(f(x))
def c(x):
    newObj = h()
    return x
class h():
    def hello_world():
        return 0"""
    )


def test_entrypoint_no_chain():
    icode = r"""Following is the code snippet:
```python
import numpy as np
from numpy import sin, cos, sum

def f(x):
    return np.sum(x)
assert f(1) == 5
def g(x):
    return cos(f(x))

def c(x):
    newObj = h()
    return x

class h():
    def hello_world():
        return 0


assert g(0) == 1
```
"""
    assert (
        sanitize(icode, "g")
        == r"""import numpy as np
from numpy import sin, cos, sum
def f(x):
    return np.sum(x)
def g(x):
    return cos(f(x))"""
    )


def test_entrypoint_variable():
    icode = r"""Following is the code snippet:
```python
import numpy as np
from numpy import sin, cos

SOME_CONSTANT = 5

def f(x):
    return c(x)
assert f(1) == 5
def g(x):
    return cos(f(x))

def c(x):
    newObj = h()
    return x

class h():
    def hello_world():
        return SOME_CONSTANT

def d(x):
    return g(x)

# Some tests
assert g(0) == 1
print(g(123))
ret = g(321)
```
"""

    assert (
        sanitize(icode, "g")
        == r"""import numpy as np
from numpy import sin, cos
SOME_CONSTANT = 5
def f(x):
    return c(x)
def g(x):
    return cos(f(x))
def c(x):
    newObj = h()
    return x
class h():
    def hello_world():
        return SOME_CONSTANT"""
    )