import matplotlib.pyplot as plt
import random


class IteratedFunctionSystem:
    def __init__(self, transforms, weights, start):
        self.transforms = transforms
        self.weights = weights
        self.start = start

    def update_position(self, x, y):
        trans, = random.choices(self.transforms, weights=self.weights)
        return trans(x, y)

    def make_trajectory(self, iterations):
        x, y = self.start
        xs, ys = [], []
        for i in range(iterations):
            xs.append(x)
            ys.append(y)
            x, y = self.update_position(x, y)
        return xs, ys

    def draw(self, iterations=100):
        xs, ys = self.make_trajectory(iterations)
        plt.scatter(xs, ys, color='green', s=1)
        plt.title(f"{iterations} iterations")
        plt.xlabel("X")
        plt.ylabel("Y")
        plt.show()


# tutaj rozne transformacje plaszczyzny
def trans_translation(x, y):
    return x + 1, y

def trans1(x, y):
    return 0.85 * x + 0.04 * y, -0.04 * x + 0.85 * y + 1.6

def trans2(x, y):
    return 0.2 * x - 0.26 * y, 0.23 * x + 0.22 * y + 1.6

def trans3(x, y):
    return -0.15 * x + 0.28 * y, 0.26 * x + 0.24 * y + 0.44

def trans4(x, y):
    return 0, 0.16 * y

def t_up(x, y):
    return x, y + 1

def t_down(x, y):
    return x, y - 1

def t_left(x, y):
    return x - 1, y

def t_right(x, y):
    return x + 1, y


def translation():
    line = IteratedFunctionSystem([trans_translation], [1], (0, 0))
    for i in [10, 100, 1000]:
        line.draw(i)


def random_walk():
    walk = IteratedFunctionSystem([t_up, t_down, t_left, t_right], [0.25, 0.25, 0.25, 0.25], (0, 0))
    for i in [100, 10000, 1000000]:
        walk.draw(i)


def barnsley():
    barnsley = IteratedFunctionSystem([trans1, trans2, trans3, trans4], [0.85, 0.07, 0.07, 0.01], (0, 0))
    for i in [100, 1000, 10000, 1000000]:
        barnsley.draw(i)


if __name__ == "__main__":
    translation()
    random_walk()
    barnsley()