#  coding: utf-8
""" RC-plot.py
    interprets a table of resistors and capacitors for impedance tokens.
    
    Table format is 8 white-space separated columns:
    
    first: label for row
    second: AB or BC    (which pair of pins used)
    six columns for impedance positions 1,2,3,4,5,6
        each column having either "--" or 
        a number followed by "ohm", "kohm", "Mohm",  "pF", "nF", "uF"

        AB = 2||(5+6)
        AC = 1
        BC = 3||4
    
    For example,

A  BC   100ohm  2.2uF   --      --      --      -- 
        means 100ohm in series with 2.2uF, to be labelled with "A"
B  BC   --      --      150ohm  100nF   --      --
        means 150ohm || 100nF, to be labelled with "B"

Plots log-log plots of |Z| vs frequency.
If label.data file exists (that is A.data, if label="A"), then
	file is read assuming frequency, phase, |Z| as columns
        and plotted on same log-log plot.

A fit is done assuming that input is of the form
	label BC  R C -- -- -- --	for series connection
        label BC  -- -- R C -- --	for parallel connection
and the fitted function is also plotted on the log-log plot.

WARNING: the assumption about series and parallel is very fragile!
"""

# To Do:
#	Make curve fitting work with fewer assumptions about
#	arrangements of components.
#
#	Change output file name to be a command-line parameter (or stdout)
#	rather than hardwired.

from __future__ import division, print_function

import sys
import matplotlib.pyplot as plt
import cmath
import numpy as np

from matplotlib.backends.backend_pdf import PdfPages
from scipy.optimize import curve_fit

from RClib import *


def good_logspace(start,stop,perdecade):
    logstart=np.log10(start)
    logstop=np.log10(stop)
    return np.logspace(logstart,logstop, np.floor((logstop-logstart)*perdecade+1.95))

freqs= good_logspace(50,5e6, perdecade=30)

multipage = PdfPages('token_plots.pdf')

for line in sys.stdin:
    if not line: continue
    row = Ztoken.from_string(line)
    absZs = [abs(row.Z(f)) for f in freqs]
    plt.loglog(freqs,absZs, label="theoretical")
    plt.title(row.to_string())
    plt.xlabel("frequency [Hz]")
    plt.ylabel("|Z| ohms")
    plt.ylim(10,100E3)
    
    zeros = [False] + [c.is_zero() for c in row.components[1:]]
    infinites = [False] + [c.is_infinite() for c in row.components[1:]]
    
    # The following is a very limited recognition of only one pattern
    # each for series and parallel configurations.
    # No check is made for more than 2 components nor for other configurations.
    is_series = row.pins=="BC" and not infinites[1] and not infinites[2]
    is_parallel = row.pins=="BC" and not infinites[3] and not infinites[4]
    
    try:
       # read the measured data (columns= freq, phase, |Z|)
       filename = row.label + ".data"
       with open(filename,'r') as datafile:
            print("DEBUG:", filename, "opened", file=sys.stderr)
            freqs=[]
            Zs = []
            for line in datafile:
               line=line.strip()
               if not line or line.startswith('#'): continue
               fields = [float(s) for s in line.split()]
               freqs.append(fields[0])
               Zs.append(fields[2])
            Zs = np.array(Zs)
#            print("DEBUG: data read, Zs=",Zs, file=sys.stderr)
            plt.loglog(freqs,Zs, label="data")
            # fit 2-componenent model 
            if is_series:
                absZf = lambda f,R,C: abs(R+ 1./(2*cmath.pi*1j*f*C))
                R = row.components[1].R
                C = row.components[2].C
                labelf = lambda R,C: to_metric(R,digits=2) + "ohm + " + to_metric(C,digits=2) + "F"
            elif is_parallel:
                absZf = lambda f,R,C: abs(Zpar(R, 1./(2*cmath.pi*1j*f*C)))
                R = row.components[3].R
                C = row.components[4].C
                labelf = lambda R,C: to_metric(R,digits=2) + "ohm || " + to_metric(C,digits=2) + "F"
#            print("DEBUG: labelf=", labelf(R,C), file=sys.stderr)
            logabsZf = lambda f,R,C: np.log(absZf(f,R,C))
            logZs = np.log(Zs)
#            print("DEBUG: len(freqs)=",len(freqs), "len(logZs)=", len(logZs), file=sys.stderr)
            fittedRC, covar = curve_fit(logabsZf, freqs, logZs, (R,C))
#            print("DEBUG: fit done", file=sys.stderr)
            plt.loglog(freqs, [absZf(f,*fittedRC) for f in freqs], label=labelf(*fittedRC))
    except IOError:
       print("Warning:", filename, "not found---ignoring", file=sys.stderr)
    plt.legend()
    multipage.savefig()
#    plt.show()
    plt.close()

multipage.close()