import sys
sys.dont_write_bytecode = True

import adcptool.sdxf as sdxf
from adcptool.vector import Vector
from adcpprocess import compute_direction_angle


def writeAscii3D(profile, format, f, voidtext="---", header=None):
    '''
    create custom ascii file that follows a user defined format 
    defined via "format" keyword, which has the following names for cell attributes
    x, y, z
    vx, vy, vz, vmag, vmag3d
    artificial
    
    '''
    output = open(f, 'w')
    
    if not header == None:
        output.write(header)
        
    for e in profile.ensembles:
        if e.void == False:
            for c in e.cells:
                output.write((format+'\n').format(
                                        x=e.position.x if hasattr(e, 'position') else voidtext,
                                        y=e.position.y if hasattr(e, 'position') else voidtext,
                                        z=c.z_position if hasattr(c, 'z_position') else voidtext,
                                        vx=c.velocity.x if hasattr(c, 'velocity') else voidtext,
                                        vy=c.velocity.y if hasattr(c, 'velocity') else voidtext,
                                        vz=c.velocity.z if hasattr(c, 'velocity') else voidtext,
                                        vmag=c.velocity.magn2d() if hasattr(c, 'velocity') else voidtext,
                                        vmag3d=c.velocity.magn() if hasattr(c, 'velocity') else voidtext,
                                        artificial=c.artificial if hasattr(c, 'artificial') else False
                                        ))
    output.close()

def writeAscii2D(profile, format, f, voidtext="---", header=None):
    '''
    create custom ascii file that follows a user defined format 
    defined via "format" keyword, which has the following names for cell attributes
    x, y, z
    vx, vy, vz, vmag, vmag3d
    depth
    ks, tau_shear, v_shear
    four_depths_1, four_depths_2, four_depths_3, four_depths_4  
    '''
    output = open(f, 'w')
    
    if not header == None:
        output.write(header)
    
    for e in profile.ensembles:
        if e.void == False:
            output.write((format+'\n').format(
                                    x=e.position.x if hasattr(e, 'position') else voidtext,
                                    y=e.position.y if hasattr(e, 'position') else voidtext,
                                    z=e.position.z if hasattr(e, 'position') else voidtext,
                                    vx=e.velocity.x if hasattr(e, 'velocity') else voidtext,
                                    vy=e.velocity.y if hasattr(e, 'velocity') else voidtext,
                                    vz=e.velocity.z if hasattr(e, 'velocity') else voidtext,
                                    vmag=e.velocity.magn2d() if hasattr(e, 'velocity') else voidtext,
                                    vmag3d=e.velocity.magn() if hasattr(e, 'velocity') else voidtext,
                                    depth=e.depth if hasattr(e, 'depth') else voidtext,
                                    four_depths_1=e.depths[0] if hasattr(e, 'depths') else voidtext,
                                    four_depths_2=e.depths[1] if hasattr(e, 'depths') else voidtext,
                                    four_depths_3=e.depths[2] if hasattr(e, 'depths') else voidtext,
                                    four_depths_4=e.depths[3] if hasattr(e, 'depths') else voidtext,
                                    ks=e.ks if hasattr(e, 'ks') else voidtext,
                                    tau_shear=e.tau_shear if hasattr(e, 'tau_shear') else voidtext,
                                    v_shear=e.v_shear if hasattr(e, 'v_shear') else voidtext
                                    ))
    output.close()    

def write2DVelocity(profile, f):
    '''
    create an ascii file with X Y and averaged velocity in U V coords
    '''

    data_2D = open(f, 'w')

    for e in profile.ensembles:
        data_2D.write('{:12.3f}{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    e.position.x, e.position.y,
                    e.velocity.x if not e.void else 0,
                    e.velocity.y if not e.void else 0))

    data_2D.close()

def writeDebug2D(profile, f):
    '''
    create an ascii file with X Y and averaged velocity in U V coords
    '''

    data_2D = open(f, 'w')

    for e in profile.ensembles:
        data_2D.write('{:12.3f}{:12.3f}{:12.3f}{:12.5f}{:12.9f}\n'.format(
                    e.position.x, e.position.y,
                    e.velocity.x if not e.void else 0,
                    e.velocity.y if not e.void else 0,
                    e.velocity.magn() if not e.void else 0,
                    compute_direction_angle(e.velocity) if not e.void else 0
                    ))

    data_2D.close()

def write2DVelocity_BlueKenue(profile, f):
    '''
    create an ascii file for Blue Kenue with adcp 2D data and node connectivity table
    '''
    # build the node connectivity table
    node1 = range(1,len(profile.ensembles))
    node2 = range(2,len(profile.ensembles) + 1)
    node3 = range(1,len(profile.ensembles))

    BK_2D = open(f, 'w')

    # header for BlueKenue
    BK_2D.write('#\n' + ':AttributeName 1 ADCP_VEL UV\n' + ':AttributeUnits 1 M/S\n'
                + '#\n' + ':NodeCount ' + str(len(profile.ensembles)) + '\n' + ':ElementCount '
                + str(len(profile.ensembles)-1) + '\n' + ':ElementType T3\n' + '#\n' + ':EndHeader\n')

    # output of adcp data and node table
    for e in profile.ensembles:
        BK_2D.write('{:12.3f}{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    e.position.x, e.position.y,
                    e.velocity.x if not e.void else 0,
                    e.velocity.y if not e.void else 0))

    for i in range(len(node1)):
        BK_2D.write(str(node1[i]) + ' ' + str(node2[i]) + ' ' + str(node3[i]) + '\n')


    BK_2D.close()

def write2DVelocity_Paraview(profile, f):
    '''
    create an ascii file for Paraview with adcp 2D data in*.vtk format
    '''

    liste_vert =range(0,len(profile.ensembles))
    liste_vert = str(liste_vert).replace(',', '').strip('[' ']')

    Paraview_2D = open(f, 'w')

    # header for BlueKenue
    Paraview_2D.write('# vtk DataFile Version 3.0\n' + 'vtk output\n' + 'ASCII\n'
                + 'DATASET POLYDATA\n' + 'POINTS ' + str(len(profile.ensembles)) + ' double\n')
    for e in profile.ensembles:
            Paraview_2D.write('{:12.3f}{:12.3f}{:12.3f}\n'.format(
                        e.position.x, e.position.y, 0.0))
    Paraview_2D.write('VERTICES 1 ' + str(len(profile.ensembles)+1) + '\n'
                + str(len(profile.ensembles)) + ' ' + liste_vert + '\n' + '\n' + 'POINT_DATA '
                + str(len(profile.ensembles)) + '\n' + 'VECTORS UV double\n')
    for e in profile.ensembles:
            Paraview_2D.write('{:12.3f}{:12.3f}{:12.3f}\n'.format(
                        e.velocity.x if not e.void else 0,
                        e.velocity.y if not e.void else 0, 0.0))

    Paraview_2D.close()

def write3DVelocity_Paraview(profile, f):
    '''
    create an ascii file for Paraview with adcp 2D data in*.vtk format
    '''

    liste_UVW = []

    for e in profile.ensembles:
        if e.void == False:
            for c in e.cells:
                liste_UVW.append(c.velocity)

    liste_vert =range(0,len(liste_UVW))
    liste_vert = str(liste_vert).replace(',', '').strip('[' ']')


    Paraview_3D = open(f, 'w')

    # header for BlueKenue
    Paraview_3D.write('# vtk DataFile Version 3.0\n' + 'vtk output\n' + 'ASCII\n'
                + 'DATASET POLYDATA\n' + 'POINTS ' + str(len(liste_UVW)) + ' double\n')

    for e in profile.ensembles:
            if e.void == False:
                for c in e.cells:
                    vel = c.velocity
                    pos = e.position + Vector(0,0, c.z_position)

                    Paraview_3D.write('{:12.3f}{:12.3f}{:12.3f}\n'.format(
                        pos.x, pos.y, pos.z))

    Paraview_3D.write('VERTICES 1 ' + str(len(liste_UVW)+1) + '\n'
                + str(len(liste_UVW)) + ' ' + liste_vert + '\n' + '\n' + 'POINT_DATA '
                + str(len(liste_UVW)) + '\n' + 'VECTORS UV double\n')

    for e in profile.ensembles:
            if e.void == False:
                for c in e.cells:
                    vel = c.velocity
                    pos = e.position + Vector(0,0, c.z_position)

                    Paraview_3D.write('{:12.3f}{:12.3f}{:12.3f}\n'.format(
                        vel.x, vel.y, vel.z))

    Paraview_3D.close()

def write3DVelocity(profile, f):
    '''
    create an ascii file whith X Y and measured velocity in U V W coords
    '''

    data_3D = open(f, 'w')

    for e in profile.ensembles:
        if e.void == False:
            for c in e.cells:
                vel = c.velocity
                pos = e.position + Vector(0,0, c.z_position)

                data_3D.write('{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    pos.x, pos.y, pos.z,
                    vel.x, vel.y, vel.z))


    data_3D.close()
    
def writeDebug3D(profile, f):
    '''
    create an ascii file whith X Y and measured velocity in U V W coords
    '''

    data_3D = open(f, 'w')

    for e in profile.ensembles:
        if e.void == False:
            for c in e.cells:
                vel = c.velocity
                pos = e.position + Vector(0,0, c.z_position)

                data_3D.write('{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    pos.x, pos.y, pos.z,
                    vel.x, vel.y, vel.z,
                    vel.magn(), compute_direction_angle(vel)
                    
                    ))


    data_3D.close()

def write3DVelocity_BlueKenue(profile, f):
    '''
    create an ascii file for Blue Kenue with adcp 3D data and node connectivity table
    '''

    liste_UVW = []

    for e in profile.ensembles:
        if e.void == False:
            for c in e.cells:
                liste_UVW.append(c.velocity)

    node1 = range(1,len(liste_UVW))
    node2 = range(2,len(liste_UVW) + 1)
    node3 = range(1,len(liste_UVW))
    node4 = range(2,len(liste_UVW) + 1)

    BK_3D = open(f, 'w')

    # header for BlueKenue
    BK_3D.write('#\n' + ':AttributeName 1 ADCP_VEL UVW\n' + ':AttributeUnits 1 M/S\n'
            + '#\n' + ':NodeCount ' + str(len(liste_UVW)) + '\n' + ':ElementCount '
            + str(len(liste_UVW)-1) + '\n' + '#\n' + ':EndHeader\n')

    # output of adcp data and node table
    for e in profile.ensembles:
        if e.void == False:
            for c in e.cells:
                vel = c.velocity
                pos = e.position + Vector(0,0, c.z_position)

                BK_3D.write('{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    pos.x, pos.y, pos.z,
                    vel.x, vel.y, vel.z))

    for i in range(len(node1)):
        BK_3D.write(str(node1[i]) + ' ' + str(node2[i]) + ' ' + str(node3[i]) + ' ' + str(node4[i]) + '\n')

    BK_3D.close()

def write3DVelocity_W0_BlueKenue(profile, f): #sets vertical velocity W = 0.0
    '''
    create an ascii file for Blue Kenue with adcp 3D data and node connectivity table
    '''

    liste_UVW = []

    for e in profile.ensembles:
        if e.void == False:
            for c in e.cells:
                liste_UVW.append(c.velocity)

    node1 = range(1,len(liste_UVW))
    node2 = range(2,len(liste_UVW) + 1)
    node3 = range(1,len(liste_UVW))
    node4 = range(2,len(liste_UVW) + 1)

    BK_3D = open(f, 'w')

    # header for BlueKenue
    BK_3D.write('#\n' + ':AttributeName 1 ADCP_VEL UVW\n' + ':AttributeUnits 1 M/S\n'
            + '#\n' + ':NodeCount ' + str(len(liste_UVW)) + '\n' + ':ElementCount '
            + str(len(liste_UVW)-1) + '\n' + '#\n' + ':EndHeader\n')

    # output of adcp data and node table
    for e in profile.ensembles:
        if e.void == False:
            for c in e.cells:
                vel = c.velocity
                pos = e.position + Vector(0,0, c.z_position)

                BK_3D.write('{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    pos.x, pos.y, pos.z,
                    vel.x, vel.y, 0.0))

    for i in range(len(node1)):
        BK_3D.write(str(node1[i]) + ' ' + str(node2[i]) + ' ' + str(node3[i]) + ' ' + str(node4[i]) + '\n')

    BK_3D.close()

def writeAverageDepth(profile, f):

    Av_depth = open(f, 'w')

    for e in profile.ensembles:
        e.depth *= -1

        Av_depth.write('{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    e.position.x, e.position.y, e.depth))

    Av_depth.close()

def write4BeamDepths(profile, f_depth4Beams, f_depthBeam1, f_depthBeam2, f_depthBeam3, f_depthBeam4):

    d_AllBeams = open(f_depth4Beams, 'w')
    d_beam1 = open(f_depthBeam1, 'w')
    d_beam2 = open(f_depthBeam2, 'w')
    d_beam3 = open(f_depthBeam3, 'w')
    d_beam4 = open(f_depthBeam4, 'w')

    for e in profile.ensembles:
        e.four_depths[0] *= -1
        e.four_depths[1] *= -1
        e.four_depths[2] *= -1
        e.four_depths[3] *= -1

        d_AllBeams.write('{:12.3f}{:12.3f}{:12.3f}{:12.3f}\n'.format(e.four_depths[0],
                    e.four_depths[1], e.four_depths[2], e.four_depths[3]))
        d_beam1.write('{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    e.position.x, e.position.y, e.four_depths[0]))
        d_beam2.write('{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    e.position.x, e.position.y, e.four_depths[1]))
        d_beam3.write('{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    e.position.x, e.position.y, e.four_depths[2]))
        d_beam4.write('{:12.3f}{:12.3f}{:12.3f}\n'.format(
                    e.position.x, e.position.y, e.four_depths[3]))

    d_AllBeams.close()
    d_beam1.close()
    d_beam2.close()
    d_beam3.close()
    d_beam4.close()

def writeDXF2D(profile, f, vel_scale=50):
    '''
    create a DXF file with the profile in 2D and its velocity vectors
    '''

    dxf = sdxf.Drawing()
    pline = [[],[],[]]

    for e in profile.ensembles:
        # convert into tuple and append
        pline[0].append(tuple(e.position))

        # draw velocity vectors too!
        if not e.void:
            dxf.append(sdxf.Line(points=[tuple(e.position), (e.position + (e.velocity * vel_scale))],
                color=1))

    dxf.append(sdxf.PolyLine(points=pline[0], closed=0, color=0))

    dxf.saveas(f)

def writeDXF3D(profile, f, vel_scale=50):
    '''
    create a DXF file with the profile in 3D and its velocity vectors
    '''

    dxf = sdxf.Drawing()

    for e in profile.ensembles:
        if not e.void:
            for c in e.cells:
                vel = vel_scale * c.velocity 
                pos = e.position + Vector(0,0, c.z_position)
                dxf.append(sdxf.Line( points=[tuple(pos), tuple(pos + vel)] ))

    dxf.saveas(f)