package com.example.polaritychecker

import android.media.AudioAttributes
import android.media.AudioFormat
import android.media.AudioTrack
import kotlinx.coroutines.*
import kotlin.math.exp
import kotlin.math.pow

class AudioGenerator {
    private var audioTrack: AudioTrack? = null
    private var isPlaying = false
    private var job: Job? = null
    private val sampleRate = 48000

    fun start(frequencyStr: String, invert: Boolean) {
        if (isPlaying) stop()
        isPlaying = true

        val frequency = when (frequencyStr) {
            "80Hz" -> 80.0
            "160Hz" -> 160.0
            "500Hz" -> 500.0
            "1k2Hz" -> 1200.0
            "4kHz" -> 4000.0
            "FR" -> 8000.0 
            else -> 1000.0
        }

        // 500ms buffer = 2 pulses per second
        val bufferSize = sampleRate / 2 
        val buffer = ShortArray(bufferSize)
        
        // Ricker wavelet centered at 50ms
        val centerTime = 0.05

        for (i in buffer.indices) {
            val t = i.toDouble() / sampleRate
            val tr = t - centerTime
            
            // Ricker wavelet equation
            val pi2f2t2 = Math.PI.pow(2) * frequency.pow(2) * tr.pow(2)
            var amplitude = (1.0 - 2.0 * pi2f2t2) * exp(-pi2f2t2)
            
            if (pi2f2t2 > 20.0) {
                amplitude = 0.0
            }
            
            if (invert) {
                amplitude = -amplitude
            }
            
            buffer[i] = (amplitude * 30000).toInt().toShort()
        }

        audioTrack = AudioTrack.Builder()
            .setAudioAttributes(
                AudioAttributes.Builder()
                    .setUsage(AudioAttributes.USAGE_MEDIA)
                    .setContentType(AudioAttributes.CONTENT_TYPE_SONIFICATION)
                    .build()
            )
            .setAudioFormat(
                AudioFormat.Builder()
                    .setEncoding(AudioFormat.ENCODING_PCM_16BIT)
                    .setSampleRate(sampleRate)
                    .setChannelMask(AudioFormat.CHANNEL_OUT_MONO)
                    .build()
            )
            .setBufferSizeInBytes(bufferSize * 2)
            .setTransferMode(AudioTrack.MODE_STREAM)
            .build()

        audioTrack?.play()

        job = CoroutineScope(Dispatchers.IO).launch {
            while (isActive && isPlaying) {
                audioTrack?.write(buffer, 0, buffer.size)
            }
        }
    }

    fun stop() {
        isPlaying = false
        job?.cancel()
        audioTrack?.pause()
        audioTrack?.flush()
        audioTrack?.release()
        audioTrack = null
    }
}
