"""
Ex. 3.7 from "A primer on...
Write a function that solves a quadratic equation, and
returns a complex number object and otherwise a float. 
To do this we need to check the value of the 
discriminant, and call cmath.sqrt if it negative and
the regular mat.sqrt otherwise. Then write two test
functions that test both cases.
"""

import math
import cmath

#a)
def roots(a, b, c):

    D = b**2 - 4 * a * c #discriminant

    if D < 0:  #solutions are complex
        x1 = (-b - cmath.sqrt(D)) / (2 * a)
        x2 = (-b + cmath.sqrt(D)) / (2 * a)
    else:      #soliutions are real
        x1 = (-b - math.sqrt(D)) / (2 * a)
        x2 = (-b + math.sqrt(D)) / (2 * a)

    return x1, x2

#b)
def test_roots_float():
    tol = 1e-14
    a, b, c = 2, 3, 1
    expected = -1.0, -0.5
    computed = roots(a, b, c)
    for e, c in zip(expected, computed):
        assert abs(e - c) < tol

def test_roots_complex():
    tol = 1e-14
    a, b, c = 1, 1, 1
    expected = (-1 - cmath.sqrt(-3)) / 2, (-1 + cmath.sqrt(-3))/ 2
    computed = roots(a, b, c)
    for e, c in zip(expected, computed):
        assert abs(e - c) < tol

test_roots_float()
test_roots_complex()

"""
Terminal> python roots_quadratic.py
(No output if the tests pass)
"""