import matplotlib.pyplot as plt
from matplotlib.ticker import AutoMinorLocator
import numpy as np
import re

# Initialize lists for Block 1 and Block 2
block1_iters = []
block1_r = []
block1_ar1 = []
block2_iters = []
block2_r = []
block2_ar1 = []
block3_iters = []
block3_r = []
block3_ar1 = []
block4_iters = []
block4_r = []
block4_ar1 = []

# Read and parse the file
with open('relative_whitening_output.txt', 'r') as file:
    for line in file:
        # Extract Iter, Block, and r values using regex
        match = re.match(r'Iter (\d+) \| Block (\d) \| x: [\d.]+ \| w: [\d.]+ \| r: ([\d.]+) \| ar1: ([\d.]+)', line)
        if match:
            iter_num = int(match.group(1))
            block_num = int(match.group(2))
            r_value = float(match.group(3))
            ar1_value = float(match.group(4))
            
            if block_num == 1:
                block1_iters.append(iter_num / 1000)  # Convert to thousands
                block1_r.append(r_value)
                block1_ar1.append(ar1_value)
            elif block_num == 2:
                block2_iters.append(iter_num / 1000)  # Convert to thousands
                block2_r.append(r_value)
                block2_ar1.append(ar1_value)

# Create the plot
plt.figure(figsize=(10, 7))
plt.plot(block1_iters, block1_r, label='First Decoder Block', linewidth=3, color='b')
plt.plot(block2_iters, block2_r, label='Second Decoder Block', linewidth=3, color='r')

# Customize the plot
plt.xlabel('Iterations (in \'000)', fontsize=20)  # Update x-axis label to indicate thousands
plt.ylabel('Ratio', fontsize=20)
# ax.yaxis.set_minor_locator(AutoMinorLocator(2))  # 4 minor intervals per major tick

xa = np.arange(0,110,20)
plt.xticks(xa, fontsize=16)
ya = np.arange(0,1.3,.2)
plt.yticks(ya,fontsize=16)
plt.title('Relative Whiteness of $W$ vs. $X$', fontsize=24)
plt.legend(fontsize=18)
plt.grid(True)
plt.ylim(bottom=0)  # Set y-axis to start at 0


# Show the plot
plt.show()
plt.ion()

# Plot the stationarity results
plt.figure(figsize=(10, 7))
plt.plot(block1_iters, block1_ar1, label='First Decoder Block', linewidth=3, color='b')
plt.plot(block2_iters, block2_ar1, label='Second Decoder Block', linewidth=3, color='r')
# plt.plot(block3_iters, block3_ar1, label='Third Decoder Block', linewidth=3)
# plt.plot(block4_iters, block4_ar1, label='Fourth Decoder Block', linewidth=3)

# Customize the plot
plt.xlabel('Iterations (in \'000)', fontsize=20)  # Update x-axis label to indicate thousands
plt.ylabel('Mean Variation', fontsize=20)
# ax.yaxis.set_minor_locator(AutoMinorLocator(2))  # 4 minor intervals per major tick

#xa = np.arange(0,110,20)
#plt.xticks(xa, fontsize=16)
#ya = np.arange(0,1.3,.2)
#plt.yticks(ya,fontsize=16)
plt.title('First-Order Stationarity', fontsize=24)
plt.legend(fontsize=14)
plt.grid(True)
plt.ylim(bottom=0)  # Set y-axis to start at 0

# Show the plot
plt.show()
