#!/usr/bin/python
import argparse
import matplotlib.pyplot as plt
import pandas as pd

fig, ((ax11, ax12), (ax21, ax22)) = plt.subplots(2, 2)

class Data:
    def __init__(self, filename):
        self.filename = filename
        self.data = {}
        self.times = {}

        i = 0
        for chunk in self.chunkYielder():
            dumpNumber, time, data = self.processChunk(chunk)
            self.data[i] = data
            self.times[i] = time
            i += 1

    def chunkYielder(self):
        lines = []
        with open(self.filename, 'r') as f:
            hasNIter = False
            hasTime = False
            hasHeaders = False
            
            # arf, not beautiful, should read it more carefully
            lines = [l.replace('\n', '') for l in f]

            i = 0

            while i < len(lines):
                niter = int(lines[i].split('#')[1])
                time = float(lines[i+1].split('#')[1])
                headers = lines[i+2].split('#')[1].split()
                data = []
                i += 3

                while i < len(lines) and '#' not in lines[i]:
                    data.append([float(val) for val in lines[i].split()])
                    i += 1
                yield niter, time, headers, data
                    
                    
                
                
    def processChunk(self, chunk):
        ''' Read the lines and get the following data in the lines:
        1: dump number
        2: time
        3: headers
        4+: data'''

        dumpNumber = chunk[0]
        time = chunk[1]
        headers = chunk[2]

        # read next line
        data = ([ [ float(el) for el in line] for line in chunk[3]])

        return dumpNumber, time, pd.DataFrame(data=data, columns=headers)

        


        
def plot(data):
    ax11.set_xdata(data['x'])
    ax11.set_ydata(data['T'])

    ax12.set_xdata(data['x'])
    ax12.set_ydata(data['S'])

    # ax21.set_xdata(data['x'])
    # ax21.set_xdata(data[''])

    ax22.set_xdata(data['T'])
    ax22.set_ydata(data['S'])

if __name__ == '__main__':
    print('Reading initial conditions')
    ic = pd.read_csv('CI.dat')
    
    print('Reading output file')
    data = Data('output.dat')

    x0 = data.data[0]

    ax11.plot(ic['x'], ic['T'])
    ax11.plot(x0['x'], x0['T'])

    ax12.plot(ic['x'], ic['S'])
    ax12.plot(x0['x'], x0['S'])

    ax22.plot(ic['T'], ic['S'])
    ax22.plot(x0['T'], x0['S'])

    plt.show()

@@ -2146,16 +2146,53 @@
 'CI.dat'
+, delim_whitespace=True)%0A    print(ic
 )%0A    %0A 
@@ -2277,16 +2277,66 @@
 ata%5B0%5D%0A%0A
+    ax11.set_xlabel('x')%0A    ax11.set_ylabel('T')%0A
     ax11
@@ -2384,32 +2384,82 @@
 'x'%5D, x0%5B'T'%5D)%0A%0A
+    ax12.set_xlabel('x')%0A    ax12.set_ylabel('S')%0A
     ax12.plot(ic
@@ -2462,32 +2462,36 @@
 t(ic%5B'x'%5D, ic%5B'S
+igma
 '%5D)%0A    ax12.plo
@@ -2504,27 +2504,92 @@
 x'%5D, x0%5B'S'%5D
+/x0%5B'x'%5D
 )%0A%0A
+    ax22.set_xlabel('$%5CSigma$')%0A    ax22.set_ylabel('T')%0A
     ax22.plo
@@ -2598,17 +2598,21 @@
 ic%5B'
-T
+Sigma
 '%5D, ic%5B'
 S'%5D)
@@ -2607,17 +2607,17 @@
 '%5D, ic%5B'
-S
+T
 '%5D)%0A    
@@ -2634,17 +2634,25 @@
 x0%5B'
-T
+S'%5D/x0%5B'x
 '%5D, x0%5B'
 S'%5D)
@@ -2647,17 +2647,17 @@
 '%5D, x0%5B'
-S
+T
 '%5D)%0A%0A   

Use Sigma