#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri May  6 16:53:45 2022

@author: maximebouchereau
"""

# Python code to drow a Barnsley fern by using dynamical systems

# Module importation

import math as mt
import numpy as np
import matplotlib.pyplot as plt
import sys
import webcolors
import datetime

# Parameters entries

A1 = [np.array([[0,0],[0,0.16]]),np.array([[0.85,0.04],[-0.04,0.85]]),np.array([[0.20,-0.26],[0.23,0.22]]),np.array([[-0.15,0.28],[0.26,0.24]])] # Matrices of the affine mappings
B1 = [np.array([0,0]),np.array([0,1.60]),np.array([0,1.60]),np.array([0,0.44])] # Vectors of the affine mappings
p1 = [0.01,0.85,0.07,0.07] # Probabilities of the iterations of each mapping
params_1 = (A1,B1,p1)

A2 = [np.array([[0,0],[0,0.25]]),np.array([[0.95,0.005],[-0.005,0.93]]),np.array([[0.035,-0.20],[0.16,0.04]]),np.array([[-0.04,0.20],[0.16,0.04]])] # Matrices of the affine mappings
B2 = [np.array([0,-0.40]),np.array([-0.002,0.50]),np.array([-0.09,0.02]),np.array([0.083,0.12])] # Vectors of the affine mappings
p2 = [0.02,0.84,0.07,0.07] # Probabilities of the iterations of each mapping
params_2 = (A2,B2,p2)

N_dots = 1000000 # Number of dot ploted over the figure

params_0 = params_2

A0, B0, p0 = A1, B1, p1

# Function for the computation of points which belong to the Barnsley fern

def Help():
    """Gives a description of the dynamical system which is used to create the Barnsley fern"""
    
    print(80*"=")
    print("    Barnsley fern: Dynamical system")
    print(80*"=")
    
    print(" ")
    print("    - Let consider:")
    print("       > Four 2x2 matrices denoted A0, A1, A2 and A3")
    print("       > Four 2x1 column vectors denoted B0, B1, B2 and B3")
    print("       > Four reals p0, p1, p2 and p3 in [0,1] such that p0+p1+p2+p3=1")
    
    print(" ")
    print("    - Let consider the sequence (x_n) of points of the plan such that:")
    print("       > x_0 = (0,0)")
    print("       > Let consider a multinomial law which gives an integer i between 0 and 3 with")
    print("       probability pi introduced in the previous paragraph")
    print("       > Let consider a rendom variable which follows the multinomial law introduced in")
    print("       the previous line, denoted X")
    print("       > The sequence (x_n) is defined by recursion with this formula:")
    print(" ")
    print("            x_{n+1} = AX x_{n} + BX")
    print(" ")
    print("       Remark: We select matrix A1 and vector B1 with probability p1, matrix A2 and vector")
    print("       B2 with probability p2, etc... and apply affine functions recursively.")
    print(" ")
    print("   - We plot the points (x_n) in the plan.")
    
    pass


def fern(N=N_dots,params=params_0):
    """
    Creates points via iterations of the mappings x -> A[i]x + B[i] with probability p[i] for 0 <= i <= 3

    Parameters
    ----------
    N : TYPE, integer
        DESCRIPTION. The default is N_dots. Number of points which will be ploted (iterations of the dynamical system)
    params : TYPE, tuple
        DESCRIPTION. The default is param_0. Parameters of the dynamical system, which contains:
        - A : TYPE, List of arrays of shape (2,2)
            DESCRIPTION. Matrices of the dynamical system
        - B : TYPE, List of arrays iof shape (2,)
            DESCRIPTION. Translations of the dynamical system
        - p : TYPE, List
            DESCRIPTION. Probabilities of choosing of matrices and vectors ofn the dynamical system

    Returns Y an array of shape (2,N)
    -------
    None.

    """
    
    A, B, p = params[0], params[1], params[2]
    Y = np.zeros((2,N))
    x = Y[:,0]
    
    print("Computation of the points...")
    for n in range(N-1):
        sys.stdout.write("\r%d   "% n+"/ "+str(N_dots))
        #sys.stdout.write("\r%d   "% int(100*(n)/(N-2))+"%")
        #sys.stdout.flush()
        idx = np.where(np.random.multinomial(1,p)==1)[0][0] # Gives the index between 0 and 3 simulated with a multinomial law af probabilities p
        x = A[idx]@x + B[idx]
        Y[:,n] = x
    
    return Y


# Function for the plot of the Barnsley fern

def B_fern(N=N_dots,params=params_0,save_fig=False,name_fig="Barnsley_fern-"+datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")):
    """
    Plots the Barnsley fern with N points, by using the mappings x -> A[i]x + B[i] for 0<= i <= 3

    Parameters
    ----------
    N : TYPE, integer
        DESCRIPTION. The default is N_dots. Number of points which will be ploted (iterations of the dynamical system)
    params : TYPE, tuple
        DESCRIPTION. The default is param_0. Parameters of the dynamical system, which contains:
        - A : TYPE, List of arrays of shape (2,2)
            DESCRIPTION. Matrices of the dynamical system
        - B : TYPE, List of arrays iof shape (2,)
            DESCRIPTION. Translations of the dynamical system
        - p : TYPE, List
            DESCRIPTION. Probabilities of choosing of matrices and vectors ofn the dynamical system
    save_fig : TYPE, boolean
             DESCRIPTION. The default is False. Saves the plot or not
    name_fig : TYPE, character string
             DESCRIPTION. Gives the nave of the saved figure (useful only if save_fig=True)

    Returns a plot of the points given by fern(N,params)
    -------
    None.

    """
    
    Y = fern(N,params)
    
    plt.figure(0)
    #colors=[webcolors.rgb_to_hex((255-int(255*i/N),int(200*i/N),0)) for i in range(N)] # Colors, depending on the order of the iteration
    #sizes = [int(np.log(i+1)) for i in range(N)]
    plt.scatter(Y[0,:],Y[1,:],color="green",marker=".",s = 2)
    plt.title("Barnsley fern")
    f = plt.gcf()
    #dpi = f.get_dpi()
    h, w = f.get_size_inches()
    f.set_size_inches(h * 0.85, w * 1.7)
    plt.axis("off")
    plt.show()
    
    if save_fig == True:
        plt.savefig(name_fig+".pdf")
    
    pass
        
        
# Have fun !