USING: kernel math math.order combinators combinators.extras generalizations math.bitwise accessors sequences ;
USING: flac.format flac.stream ;
USING: prettyprint ;

IN: flac.decoder

SYMBOL: current-flac-output

: decode-block-size ( n -- n )
    dup
    {
        { [ 0b0000 = ]               [ drop reserved-block-size ] }
        { [ 0b0001 = ]               [ drop 192 ] }
        { [ 0b0010 0b0101 between? ] [ 2 - 2^ 567 * ] }
        { [ 0b0110 = ]               [ drop 8 flac-read 1 + ] }
        { [ 0b0111 = ]               [ drop 16 flac-read 1 + ] }
        { [ 0b1000 0b1111 between? ] [ 8 - 2^ 256 * ] }
    } cond-case ;

: decode-sample-rate ( n -- n )
    {
        { 0b0000 [ "TODO: GET FROM STREAMINFO" ] }
        { 0b0001 [ 88200 ] }
        { 0b0010 [ 17640 ] }
        { 0b0011 [ 19200 ] }
        { 0b0100 [ 8000 ] }
        { 0b0101 [ 16000 ] }
        { 0b0110 [ 22050 ] }
        { 0b0111 [ 24000 ] }
        { 0b1000 [ 32000 ] }
        { 0b1001 [ 44100 ] }
        { 0b1010 [ 48000 ] }
        { 0b1011 [ 96000 ] }
        { 0b1100 [ 8 flac-read 1000 * ] } ! sample rate in kHz
        { 0b1101 [ 16 flac-read ] }       ! sample rate in Hz
        { 0b1110 [ 16 flac-read 10 * ] }  ! sample rate in tens of Hz
        { 0b1111 [ invalid-sample-rate ] }
    } case ;

: decode-channel-assignment ( n -- channel-assignment )
    {
        { [ 0b0000 0b0111 between? ] [ 0 ] }
        { [ 0b1000 = ]               [ 1 ] }
        { [ 0b1001 = ]               [ 2 ] }
        { [ 0b1010 = ]               [ 3 ] }
        [ reserved-channel-assignment ]
    } cond-case <flac-channel-assignment> ;

: decode-bits-per-sample ( stream-info n -- n )
    {
        { 0b000 [ 0 ] }
        { 0b001 [ 8 ] }
        { 0b010 [ 12 ] }
        { 0b011 [ reserved-sample-size ] }
        { 0b100 [ 16 ] }
        { 0b101 [ 20 ] }
        { 0b110 [ 24 ] }
        { 0b111 [ reserved-sample-size ] }
    } case
    dup 0 = [ drop bits-per-sample>> ] [ nip ] if ;

: read/assert-frame-sync-code ( -- )
    15 flac-read 0b111111111111100 = [ sync-code-error ] unless ;

: read-flac-frame-header ( stream-info -- frame-header )
    [
        read/assert-frame-sync-code
        1 flac-read <flac-frame-blocking-strategy>
        4 flac-read ! blocksize code
        4 flac-read ! sample rate code
        4 flac-read decode-channel-assignment
    ] dip
    3 flac-read decode-bits-per-sample
    1 flac-read drop ! ignore magic sync for now
    flac-read-coded-number
    ! decoding blocksize/samplerate potentially reads from end of header
    [ decode-block-size ] 4dip
    [ decode-sample-rate ] 3dip
    8 flac-read
    flac-frame-header boa ;

: read-flac-subframe-wasted-bits ( -- k )
    1 flac-read 1 = [ 0 [ 1 + 1 flac-read 0 = ] loop ] [ 0 ] if ;

: read-flac-subframe-constant ( blocksize bps wasted-bits -- constant-subframe )
    [ flac-read-int32 ] dip shift <repetition> [  ] map flac-subframe-constant boa ;

: read-flac-subframe-verbatim ( blocksize bps wasted-bits -- verbatim-subframe )
    [ <repetition> ] dip swap [ flac-read-int32 swap shift ] with map flac-subframe-verbatim boa ;

: read-residual-coding-method-type ( -- coding-method param-bits param-escape )
    2 flac-read dup
    2 > [ reserved-residual-coding-type ] when
    <flac-entropy-coding-method-type>
    dup entropy-coding-partitioned-rice =  [ 4 0xf ] [ 5 0x1f ] if ;

: read-flac-subframe-lpc ( frame-header subframe-header -- lpc-subframe )
    2drop 1 ;

:: flac-read-residuals ( blocksize -- sequence )
    read-residual-coding-method-type :> ( coding-type param-bits param-escape )
    1 ;

: read-flac-subframe-warmup-samples ( bps order -- seq )
      <repetition> [ flac-read-int32 ] map ;

:: read-flac-subframe-fixed (  blocksize bps order -- fixed-subframe )
    bps order read-flac-subframe-warmup-samples :> warmup
    4 flac-read :> partition-order
    partition-order 2^ :> partitions

!    2dup [ read-flac-subframe-warmup-samples ] [ flac-read-residuals ] 2bi*
!    flac-subframe-fixed boa ;

: read-flac-subframe-qlp-coeff-precision ( -- precision )
    4 flac-read 1 + ;

: read-flac-qlp-shift ( -- quantization-level )
    5 flac-read ;

: read-flac-qlp-coeff ( subframe-header coeff-precision -- seq )
    [ predictor-order>> ] dip <repetition> [ flac-read ] map ;

! 000000 constant
! 000001 verbatim
! 00001x reserved
! 0001xx reserved
! 001xxx if xxx <= 4 fixed, xxx = order ; else reserved
! 01xxxx reserved
! 1xxxxx lpc
: read-flac-subframe-type ( -- subframe-type order )
    6 flac-read dup
    {
        { [ 0b0000 = ] [ drop f subframe-type-constant ] }
        { [ 0b0001 = ] [ drop f subframe-type-verbatim ] }
        { [ 0b1000 0b1100 between? ] [ 3 clear-bit subframe-type-fixed ] }
        { [ 0b100000 0b111111  between? ] [ 5 clear-bit 1 + subframe-type-lpc ] }
        [ drop reserved-subframe-type ]
    } cond-case swap ;

: read/assert-flac-subframe-sync ( -- )
    1 flac-read 1 = [ invalid-subframe-sync ] when ;

: read-flac-subframe-header ( -- subframe-header )
    read/assert-flac-subframe-sync
    read-flac-subframe-type
    read-flac-subframe-wasted-bits
    flac-subframe-header boa ;

: calculate-bits-per-sample ( bps wasted-bits subframe-type channel -- bps' )
    [ - ] 2dip swap
    {
        { channel-assignment-left  [ 1 = [ 1 ] [ 0 ] if ] }
        { channel-assignment-right [ 0 = [ 1 ] [ 0 ] if ] }
        { channel-assignment-mid   [ 1 = [ 1 ] [ 0 ] if ] }
        [ 2drop 0 ]
    } case + ;

:: read-flac-subframe ( frame-header channel -- subframe )
    read-flac-subframe-header :> subframe-header
    subframe-header wasted-bits>> :> wasted-bits
    subframe-header subframe-type>> :> subframe-type
    subframe-header predictor-order>> :> order
    frame-header bits-per-sample>> wasted-bits subframe-type channel calculate-bits-per-sample :> bps
    frame-header blocksize>> :> blocksize

    subframe-type
    {
        { subframe-type-constant [ blocksize bps wasted-bits read-flac-subframe-constant ] }
        { subframe-type-verbatim [ blocksize bps wasted-bits read-flac-subframe-verbatim ] }
!        { subframe-type-fixed    [ blocksize bps wasted-bits read-flac-subframe-fixed    ] }
!        { subframe-type-lpc      [ blocksize bps wasted-bits read-flac-subframe-lpc      ] }
    } case
    subframe-header swap flac-subframe boa ;

: read-flac-subframes ( stream-info frame-header -- seq )
    [ channels>> ] dip swap <iota> [ read-flac-subframe ] with map ;

: read-flac-frame-footer ( -- frame-footer )
    16 flac-read flac-frame-footer boa ;

: read-flac-frame ( stream-info -- frame )
    [ read-flac-frame-header ] keep swap
    [ read-flac-subframes ] keep swap
    read-flac-frame-footer
    flac-frame boa ;