#! /bin/env python3
# -*- coding: utf-8 -*-

#
#   Simulate open loop static crossover distortion of EF output stage
#   Output distortion into WAV file so it can be listened to.
#
#
#   Caveat: this uses a simple exponential BJT plus emitter resistor model:
#
#       No drivers, no beta droop, etc
#       No switching artifacts
#       No charge storage
#       No power supply
#       No feedback
#

import time, subprocess, sys, pprint, hashlib
from multiprocessing import Pool
import numpy as np
import scipy.interpolate
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter
from path import Path

TMPDIR = Path("/mnt/ssd/temp")
CHUNK_SIZE  = 131072    # in samples
N_CORES = 6             # number of processes (not counting sox)

#
#   Output parameters
#
PEAK_VOLTAGE = 12               # Volume control setting for digital full scale
LOUDSPEAKER_IMPEDANCE = 4.0     # Ohms resistive load

class Const(object):
    # physical constants
    q = 1.6021766208e-19  # electron charge
    k = 1.38064852e-23    # boltzmann constant
    T = 300               # temperature, kelvins

    Vt=k*T/q

def dB( x ):
    return 10*np.log10( (x*x.conj()).real )

#   Solves (and precalculates) the BJT equations for the following circuit :
#
#         VCC
#          |
#   Vin--|<
#          |
#          R
#          |
#         GND
#
#   Parameter "R" is the emitter resistor.
#   Ie = I(R) will span from 0 to parameter "imax", using "steps" values.
#   Early effect etc is not modeled.
#
def generate_transistor_iv( R, imax, steps, IS=9.8145e-12 ):
    Vt = Const.Vt
    I = np.logspace( np.log10(imax*1e-9), np.log10(imax), steps, True, dtype=np.float64 )[1:]
    Vin = R*I + Vt*np.log( I/IS )
    return I, Vin

#
#   models a class-AB common collector push pull BJT output stage
#   Re1, IS1 : top NPN
#   Re2, IS1 : bottom PNP
#
class OutputStage( object ):
    def __init__( self, Re, Rs=0., IS1=5E-12, IS2=7e-12 ):
        self.Re1 = Re
        self.Re2 = Re
        self.Rs  = Rs   # series resistor in the output
        self.IS1 = IS1
        self.IS2 = IS2

    #   Precalculate distortion tables for :
    #   ibias : idle bias current
    #   imax  : maximum current to compute
    #   steps : number of simulation points
    def set_bias( self, ibias, imax, steps ):
        self.imax = imax
        self.ibias = ibias

        # generate characteristics for each transistor and its associated resistor
        I1, V1 = generate_transistor_iv( self.Re1, imax+ibias+1e-3, steps, self.IS1 )
        I2, V2 = generate_transistor_iv( self.Re2, imax+ibias+1e-3, steps, self.IS2 )

        # compute bias voltage : find Vin for our common emitter transistor
        # which gives the requested bias current. This gives each transistor's
        # bias voltage. The voltage on the Vbe multiplier bias spreader would be the sum.
        VBias1, = np.interp( [ibias], I1, V1 )
        VBias2, = np.interp( [ibias], I2, V2 )

        # Add bias voltage to both "Vbe+V(Re)" vectors
        # So when Vin=Vout, Iout=0 and I1=I2=ibias
        V1 -= VBias1
        V2  = -(V2 - VBias2)    # also flip the PNP's voltage

        # From Vd = Vout-Vin, calculate current in each transistor
        # and resample it so we have a precalculated vector using Vd as axis.
        Vd = np.unique( np.hstack(( V1, V2 )))  # sort and remove duplicates
        I1 = np.interp( Vd, V1, I1 )
        I2 = np.interp( Vd, V2[::-1], I2[::-1] )    # flip arrays, interp needs increasing order

        # output current into load
        Iout = I1 - I2

        # snip off the invalid excess values
        mi, ma = np.searchsorted( Iout, (-imax, imax) )
        mi = max( mi-1, 0 ) # ensure [-imax, imax] interval is included inside returned values

        # Remove linear output impedance, keep only distortion
        # If gm-doubling occurs, this will invert distortion polarity in the part of the
        # waveform through the crossover when  output impedance is lower than Re. So it's
        # a bit wonky.
        # Vd -= Iout*self.Re1

        # Vd = (Vout-Vin) is the distortion added by the output stage, it depends on Iout
        self._disto_Iout = Iout[mi:ma]
        self._disto_Vd   = Vd[mi:ma]
        self._disto_Iout_min = self._disto_Iout[0]
        self._disto_Iout_max = self._disto_Iout[-1]

        return self._disto_Iout, self._disto_Vd

    #   calculate distortion due to output stage when Iout = "current"
    def get_error_voltage( self, current ):
        assert current.max() <= self._disto_Iout_max
        assert current.min() >= self._disto_Iout_min
        return np.interp( current, self._disto_Iout, self._disto_Vd )

    # Helper : Plot a wing diagram
    # ibias should be an array of bias currents to plot
    def plot_wing_diagram( self, ibias_list=[None], imax=None, steps=None ):
        # plt.figure()
        for ibias in ibias_list:
            if ibias != None:
                self.set_bias( ibias, imax, steps )
            i, v = self._disto_Iout, self._disto_Vd
            label = "%d mA %.02fR"%(self.ibias*1000,self.Re1)
            y = np.diff(v)/np.diff(i)
            plt.plot( i[1:], y, label=label)
            xpos = int(len(i)*0.6)
            if self.Re1:
                xpos=0
            plt.text( i[xpos], y[xpos], label )
        plt.xlabel( "Iout (A)" )
        plt.ylabel( "dV/dI = dynamic resistance" )
        plt.grid( True )
        # plt.legend()


# Has to be a class for multiprocessing
class DistortFile( object ):
    def __init__( self, _fname, _ops_list ):
        self.fname = Path( _fname )
        self.ops_list = _ops_list

        # get file info
        print( '\nFile "%s"' % self.fname )
        p = subprocess.run([ "soxi", self.fname ], capture_output=True, encoding=sys.stdout.encoding )
        attrs = {}
        for line in p.stdout.split("\n"):
            kv = line.split(":")
            if len(kv)==2:
                attrs[kv[0].strip().lower()] = kv[1].strip()

        pprint.pprint( attrs )
        self.srate_file  = int(attrs["sample rate"])
        self.channels    = int(attrs["channels"])

        # get base sample rate
        if not self.srate_file % 44100:
            self.srate_base = 44100
        elif not srate_file % 48000:
            self.srate_base = 48000
        else:
            raise ValueError("Only 48k and 44.1k base rates supported.")

        # decide sample rates for oversampling
        self.srate_high = max( self.srate_file, self.srate_base*64 )
        print( "Oversampling %.02fk -> %.02fk" % ( self.srate_file/1000, self.srate_high/1000 ))

        # oversample file and store in temporary file
        # to avoid doing it again if we want to change the distortion settings
        self.tmpfile = TMPDIR/hashlib.sha256(self.fname.encode("utf-8")).hexdigest()+".raw"

        if not self.tmpfile.isfile():
            cmd = [ "sox", self.fname, "--encoding", "float", "--bits", "32",
                    "--type", "raw", "--endian", "little", "--no-dither",
                    str(self.tmpfile), "rate", "-v", str(self.srate_high) ]
            print( cmd )
            subprocess.run( cmd )
        self.raw_size = self.tmpfile.size
        print( "Oversampled data %.02f Gb" % (self.raw_size/1024**3) )

        # Parallel processing
        with Pool( N_CORES ) as p:
            p.map( self.proc, self.ops_list )

    def proc( self, ops ):
        # Setup output file
        outfile = TMPDIR/self.fname.basename()+(" %.03fR %.03fR %.1fmA.wav"%(ops.Re1,ops.Re2,ops.ibias) )
        print( "Writing", outfile )

        # Use sox to downsample our high sample rate simulation back to original sample rate
        cmd = [ "sox", "--encoding", "float", "--bits", "32", "--type", "raw", "--endian", "little", "--rate", str(self.srate_high), "--channels", str(self.channels), "-",
                "--encoding", "signed", "--bits", "24", str(outfile), "rate", "-v", str(self.srate_file) ]
        p = subprocess.Popen( cmd, stdin=subprocess.PIPE )

        # mmap input file
        data = np.memmap( self.tmpfile, dtype=np.float32, mode='r', offset=0, shape=(self.channels, self.raw_size//(4*self.channels)), order='F')

        # process it in chunks, can't load the whole file in memory
        t = 0
        for pos in range( 0, data.shape[1], CHUNK_SIZE ):
            chunk = data[:, pos:(pos+CHUNK_SIZE) ] * (PEAK_VOLTAGE/LOUDSPEAKER_IMPEDANCE)   # get output current
            disto = ops.get_error_voltage( chunk ) * (1.0/PEAK_VOLTAGE)     # get distortion and normalize it to input signal
            p.stdin.write( disto.astype(np.float32).tobytes(order="F") )    # send to sox
            tt = time.time()
            if tt>t:
                sys.stdout.write( "%.1f%%\r" % (pos*100/data.shape[1]))
                t = tt+1.0

if __name__ == '__main__':

    plt.rc('font', size=12)
    ops_list = []

    # Pairs of (Bias in amps, Re in ohms)
    for bias,Re in [(0.02, 0.51), (0.05,0.47), (0.1, 0.22), (0.2, 0.1)]:

        # do it without emitter resistors
        ops = OutputStage( Re=0, Rs=0 )
        ops.set_bias( bias, 1.1*PEAK_VOLTAGE/LOUDSPEAKER_IMPEDANCE, 65536 )
        ops_list.append( ops )

        # do it with emitter resistors
        ops = OutputStage( Re=Re )
        ops.set_bias( bias, 1.1*PEAK_VOLTAGE/LOUDSPEAKER_IMPEDANCE, 65536 )
        ops_list.append( ops )

    # plot it before running it
    for ops in ops_list:
        ops.plot_wing_diagram()
    plt.show()

    DistortFile(  "/mnt/disk/music/flac/albums/Rebecca Pidgeon/The Raven/05 Grandmother.flac", ops_list )