
import matplotlib.pyplot as plt


def plot_guess(x, guess, smallest_guess_too_big, largest_guess_too_small):
    height = .5
    line_height = height / 2
    xmin = 0
    xmax = x
    fig = plt.figure(figsize=(10, height))
    ax = fig.add_subplot(111)
    ax.set_ylim(0, height)

    # draw horizontal line
    plt.hlines(line_height, xmin, xmax)
    
    # draw vertical lines
    def plot_vline(x_val, color=None):
        plt.vlines(x_val, 0, height, colors=color)
    plot_vline(xmin)
    plot_vline(xmax)
    plot_vline(smallest_guess_too_big, "red")
    plot_vline(largest_guess_too_small, "red")
    

    # draw a point on the line
    plt.plot(guess, line_height, 'ro', ms = 4, mfc = 'r')

    # add bookend and guess
    plt.text(xmin - 0.1, height/8*3, xmin, horizontalalignment='right')
    plt.text(xmax + 0.1, height/8*3, xmax, horizontalalignment='left')

    plt.axis('off')
    plt.show()
    

def nth_root(x: float, n: int):
    """Calculate the nth root of x"""
    
    if x < 0:
        raise NotImplemented("not implemented for negative values")
    if n < 0:
        raise NotImplemented("not implemented for negative root values")
        
    if x == 0 or x == 1:
        return x
    
    guess = x / 2
    prev_guess = None
    
    if x > 1:
        smallest_guess_too_big = x
        largest_guess_too_small = 0
    else:  # x < 1
        smallest_guess_too_big = 1
        largest_guess_too_small = x
    
    guess_nth_power = guess ** n 
    
    while (guess_nth_power != x and guess != prev_guess):
        plot_guess(x, guess, smallest_guess_too_big, largest_guess_too_small)
        text = f"guess: {guess}^{n} = {guess_nth_power}"
        if guess_nth_power > x:
            smallest_guess_too_big = guess
            prev_guess = guess
            new_guess = (largest_guess_too_small + guess) / 2
            text = f"{text} > {x}"
        else:
            largest_guess_too_small = guess
            prev_guess = guess
            new_guess = (guess + smallest_guess_too_big) / 2
            text = f"{text} < {x}"
            
        print(text)
        guess_nth_power = new_guess ** n
        guess = new_guess
        
    plot_guess(x, guess, smallest_guess_too_big, largest_guess_too_small)
    print(f"{guess}^{n} = {guess_nth_power}, woo!")
        
    return guess


nth_root(8, 2)

guess: 4.0^2 = 16.0 > 8

guess: 2.0^2 = 4.0 < 8

guess: 3.0^2 = 9.0 > 8

guess: 2.5^2 = 6.25 < 8

guess: 2.75^2 = 7.5625 < 8

guess: 2.875^2 = 8.265625 > 8

guess: 2.8125^2 = 7.91015625 < 8

guess: 2.84375^2 = 8.0869140625 > 8

guess: 2.828125^2 = 7.998291015625 < 8

guess: 2.8359375^2 = 8.04254150390625 > 8


nth_root(16, 2)

guess: 8.0^2 = 64.0 > 16

4.0^2 = 16.0, woo!

4.0


nth_root(625, 4)

guess: 312.5^4 = 9536743164.0625 > 625

guess: 156.25^4 = 596046447.7539062 > 625

guess: 78.125^4 = 37252902.98461914 > 625

guess: 39.0625^4 = 2328306.4365386963 > 625

guess: 19.53125^4 = 145519.15228366852 > 625

guess: 9.765625^4 = 9094.947017729282 > 625

guess: 4.8828125^4 = 568.4341886080801 < 625

guess: 7.32421875^4 = 2877.6980798284058 > 625

guess: 6.103515625^4 = 1387.7787807814457 > 625

guess: 5.4931640625^4 = 910.5216580707065 > 625


nth_root(0.5, 2)

guess: 0.25^2 = 0.0625 < 0.5

guess: 0.625^2 = 0.390625 < 0.5

guess: 0.8125^2 = 0.66015625 > 0.5

guess: 0.71875^2 = 0.5166015625 > 0.5

guess: 0.671875^2 = 0.451416015625 < 0.5

guess: 0.6953125^2 = 0.48345947265625 < 0.5

guess: 0.70703125^2 = 0.4998931884765625 < 0.5

guess: 0.712890625^2 = 0.5082130432128906 > 0.5

guess: 0.7099609375^2 = 0.5040445327758789 > 0.5

guess: 0.70849609375^2 = 0.5019667148590088 > 0.5